ensure vthr is numpy array, not scalar

jeshraghian · Jun 20, 2024 · 49a4b35 · 49a4b35
1 parent 2cb5bcb
commit 49a4b35
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/snntorch/export_nir.py b/snntorch/export_nir.py
@@ -27,6 +27,7 @@ def _extract_snntorch_module(module: torch.nn.Module) -> Optional[nir.NIRNode]:
 
         beta = module.beta.detach().numpy()
         vthr = module.threshold.detach().numpy()
+        vthr = np.array([vthr]) if isinstance(vthr, (int, float)) else vthr
         tau_mem = dt / (1 - beta)
         r = tau_mem / dt
         v_leak = np.zeros_like(beta)
@@ -56,6 +57,7 @@ def _extract_snntorch_module(module: torch.nn.Module) -> Optional[nir.NIRNode]:
         alpha = module.alpha.detach().numpy()
         beta = module.beta.detach().numpy()
         vthr = module.threshold.detach().numpy()
+        vthr = np.array([vthr]) if isinstance(vthr, (int, float)) else vthr
 
         tau_syn = dt / (1 - alpha)
         tau_mem = dt / (1 - beta)