Merge pull request #211 from kklurz/main

Minor bug fixes for zero_inflation losses

neuralpredictors/layers/encoders/base.py

@@ -1,3 +1,4 @@
+import torch
 from torch import nn
 
 
@@ -73,6 +74,12 @@ def forward(
         else:
             x = self.readout(x, data_key=data_key, shift=shift)
 
+        # keep batch dimension if only one image was passed
+        params = []
+        for param in x:
+            params.append(param[None, ...] if len(param.shape) == 1 else param)
+        x = torch.stack(params)
+
         if self.modulator:
             x = self.modulator[data_key](x, behavior=behavior)
 

neuralpredictors/layers/encoders/firing_rate.py

@@ -73,6 +73,7 @@ def forward(
             shift = self.shifter[data_key](pupil_center, trial_idx)
 
         x = self.readout(x, data_key=data_key, shift=shift, **kwargs)
+        x = x[None, ...] if len(x.shape) == 1 else x  # keep dimensions if only one image was passed
 
         if self.modulator:
             if behavior is None:

neuralpredictors/layers/encoders/zero_inflation_encoders.py

@@ -103,6 +103,12 @@ def forward_base(
         else:
             x = self.readout(x, data_key=data_key, shift=shift)
 
+        # keep batch dimension if only one image was passed
+        params = []
+        for param in x:
+            params.append(param[None, ...] if len(param.shape) == 1 else param)
+        x = torch.stack(params)
+
         if self.modulator:
             x = self.modulator[data_key](x, behavior=behavior)
 

neuralpredictors/measures/zero_inflated_losses.py

@@ -17,8 +17,8 @@ def forward(self, target, output, **kwargs):
         if loc.requires_grad:
             self.multi_clamp(loc, [0.0] * neurons_n, target.max(dim=0)[0])
 
-        zero_mask = target < loc
-        nonzero_mask = target >= loc
+        zero_mask = target <= loc
+        nonzero_mask = target > loc
 
         # spike loss
         spike_logl = torch.log(1 - q) - torch.log(loc)