DiscreteCurve.forward: drop batch dim if it is 1

Author: sorenhauberg

stochman/curves.py

@@ -199,6 +199,8 @@ def forward(self, t: torch.Tensor) -> torch.Tensor:
             torch.floor(tt * num_edges).clamp(min=0, max=num_edges - 1).long()  # Bx|t|
         ).unsqueeze(2).repeat(1, 1, D)  # Bx|t|xD, this assumes that nodes are equi-distant
         result = torch.gather(a, 1, idx) * tt.unsqueeze(2) + torch.gather(b, 1, idx)  # Bx|t|xD
+        if B == 1:
+            result = result.squeeze(0)  # |t|xD
         return result
 
     def __getitem__(self, indices: int) -> "DiscreteCurve":