#488 fix dimensions issues

#488 fix dimensions issues with layers with different shape
BindsNET · May 19, 2021 · 159cabe · 159cabe
1 parent ea593ad
commit 159cabe
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/bindsnet/learning/learning.py b/bindsnet/learning/learning.py
@@ -552,11 +552,13 @@ 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, device=self.source.s.device
+                #batch_size, *self.source.shape, device=self.source.s.device
+                batch_size, self.source.n, device=self.source.s.device
             )
         if not hasattr(self, "p_minus"):
             self.p_minus = torch.zeros(
-                batch_size, *self.target.shape, device=self.target.s.device
+                # batch_size, *self.target.shape, device=self.target.s.device
+                batch_size, self.target.n, device=self.target.s.device
             )
         if not hasattr(self, "eligibility"):
             self.eligibility = torch.zeros(