device fix

SpikeAI · Nov 19, 2023 · 284ce89 · 284ce89
1 parent 0a42690
commit 284ce89
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/snntorch/_neurons/leakyparallel.py b/snntorch/_neurons/leakyparallel.py
@@ -252,8 +252,9 @@ def _diagonal_enable(self, diagonal_enable):
                         # self.rnn.weight_hh_l0.grad[i, j] = 0
 
     def grad_hook(self, grad):
+        device = grad.device
         # Create a mask that is 1 on the diagonal and 0 elsewhere
-        mask = torch.eye(self.hidden_size, self.hidden_size)
+        mask = torch.eye(self.hidden_size, self.hidden_size, device=device)
         # Use the mask to zero out non-diagonal elements of the gradient
         return grad * mask