Ensures single precision on GPU

ORNL · Sep 27, 2023 · 183b878 · 183b878
1 parent 529a268
commit 183b878
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/superneuroabm/core/util.py b/superneuroabm/core/util.py
@@ -25,7 +25,7 @@ def find_max_depth(l, curr_depth=0):
         print("no dims specified... running find max depth")
         find_max_depth(tensor)
         max_dims = tuple(list(dim2maxlen.values()))
-    answer = np.full(shape=max_dims, fill_value=np.nan)
+    answer = np.full(shape=max_dims, fill_value=np.nan, dtype=np.float32)
 
     def fill_arr(arr, coord):
         if len(coord) == len(max_dims):

diff --git a/superneuroabm/neuron.py b/superneuroabm/neuron.py
@@ -98,9 +98,10 @@ def synapse_step_func(
     # Update outgoing synapses if any
     for synapse_idx in range(len(output_synapsess[my_idx])):
         out_synapse_info = output_synapsess[my_idx][synapse_idx]
-        out_neuron_id = int(out_synapse_info[0])
+        out_neuron_id = out_synapse_info[0]
         if math.isnan(out_neuron_id):
             break
+        out_neuron_id = int(out_neuron_id)
         # If out_neuron still in refractory period, return
         if t_elapses[out_neuron_id] > 0:
             continue
@@ -150,6 +151,7 @@ def synapse_with_stdp_step_func(
         out_neuron_id = out_synapse_info[0]
         if math.isnan(out_neuron_id):
             break
+        out_neuron_id = int(out_neuron_id)
         weight = out_synapse_info[1]
         synapse_register = out_synapse_info[2:]
         # Check if delayed Vm was over threshold, if so spike