Merge pull request BindsNET#406 from BindsNET/hananel

Move reward modulation inputs to GPU

Author: Hananel-Hazan

bindsnet/learning/learning.py

@@ -13,6 +13,7 @@
 )
 from ..utils import im2col_indices
 
+
 class LearningRule(ABC):
     # language=rst
     """
@@ -547,11 +548,17 @@ def _connection_update(self, **kwargs) -> None:
 
         # Initialize eligibility, P^+, and P^-.
         if not hasattr(self, "p_plus"):
-            self.p_plus = torch.zeros(batch_size, *self.source.shape)
+            self.p_plus = torch.zeros(
+                batch_size, *self.source.shape, device=self.source.s.device
+            )
         if not hasattr(self, "p_minus"):
-            self.p_minus = torch.zeros(batch_size, *self.target.shape)
+            self.p_minus = torch.zeros(
+                batch_size, *self.target.shape, device=self.target.s.device
+            )
         if not hasattr(self, "eligibility"):
-            self.eligibility = torch.zeros(batch_size, *self.connection.w.shape)
+            self.eligibility = torch.zeros(
+                batch_size, *self.connection.w.shape, device=self.connection.w.device
+            )
 
         # Reshape pre- and post-synaptic spikes.
         source_s = self.source.s.view(batch_size, -1).float()