Update label_tensor.py cpu/gpu (#292)

* Update label_tensor.py cpu/gpu
* Update test_adaptive_refinment_callbacks.py
* Update test_optimizer_callbacks.py

pina/label_tensor.py

@@ -176,7 +176,7 @@ class LabelTensor(torch.Tensor):
         tmp = super().cuda(*args, **kwargs)
         new = self.__class__.clone(self)
         new.data = tmp.data
-        return tmp
+        return new
 
     def cpu(self, *args, **kwargs):
         """
@@ -185,7 +185,7 @@ class LabelTensor(torch.Tensor):
         tmp = super().cpu(*args, **kwargs)
         new = self.__class__.clone(self)
         new.data = tmp.data
-        return tmp
+        return new
 
     def extract(self, label_to_extract):
         """

tests/test_callbacks/test_adaptive_refinment_callbacks.py

@@ -71,6 +71,7 @@ def test_r3refinment_routine():
     # make the trainer
     trainer = Trainer(solver=solver,
                       callbacks=[R3Refinement(sample_every=1)],
+                      accelerator='cpu',
                       max_epochs=5)
     trainer.train()
 

tests/test_callbacks/test_optimizer_callbacks.py

@@ -84,5 +84,6 @@ def test_switch_optimizer_routine():
                                           new_optimizers_kwargs={'lr': 0.01},
                                           epoch_switch=3)
                       ],
+                      accelerator='cpu',
                       max_epochs=5)
     trainer.train()