Fix bug in handling labels with LabelTensor (#460)

---------

Co-authored-by: Filippo Olivo <[email protected]>

Author: MatteB03

pina/label_tensor.py

@@ -448,6 +448,9 @@ def __getitem__(self, index):
 
         # Retrieve selected tensor and labels
         selected_tensor = super().__getitem__(index)
+        if not hasattr(self, "_labels"):
+            return selected_tensor
+
         original_labels = self._labels
         updated_labels = copy(original_labels)
 

pina/model/block/convolution_2d.py

@@ -265,7 +265,7 @@ def _make_grid_transpose(self, X):
 
         """
         # initialize to all zeros
-        tmp = torch.zeros_like(X)
+        tmp = torch.zeros_like(X).as_subclass(torch.Tensor)
         tmp[..., :-1] = X[..., :-1]
 
         # save on tmp