fixing cuda/cpu

mathLab · Oct 18, 2023 · c8760d2 · c8760d2
1 parent 1f0e8c7
commit c8760d2
Showing 1 changed file with 14 additions and 12 deletions.
diff --git a/pina/label_tensor.py b/pina/label_tensor.py
@@ -132,20 +132,22 @@ def select(self, *args, **kwargs):
         return tmp
 
     def cuda(self, *args, **kwargs):
-            """
-            Send Tensor to cuda. For more details, see :meth:`torch.Tensor.cuda`.
-            """
-            tmp = super().cuda(*args, **kwargs)
-            tmp._labels = self._labels
-            return tmp
+        """
+        Send Tensor to cuda. For more details, see :meth:`torch.Tensor.cuda`.
+        """
+        tmp = super().cuda(*args, **kwargs)
+        new = self.__class__.clone(self)
+        new.data = tmp.data
+        return tmp
 
     def cpu(self, *args, **kwargs):
-            """
-            Send Tensor to cpu. For more details, see :meth:`torch.Tensor.cpu`.
-            """
-            tmp = super().cpu(*args, **kwargs)
-            tmp._labels = self._labels
-            return tmp
+        """
+        Send Tensor to cpu. For more details, see :meth:`torch.Tensor.cpu`.
+        """
+        tmp = super().cpu(*args, **kwargs)
+        new = self.__class__.clone(self)
+        new.data = tmp.data
+        return tmp
 
     def extract(self, label_to_extract):
         """