fixing tests

mathLab · Nov 7, 2023 · 4023eb0 · 4023eb0
1 parent 05050f4
commit 4023eb0
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/pina/callbacks/adaptive_refinment_callbacks.py b/pina/callbacks/adaptive_refinment_callbacks.py
@@ -61,7 +61,7 @@ def _compute_residual(self, trainer):
             pts.retain_grad()
             # PINN loss: equation evaluated only on locations where sampling is needed
             target = condition.equation.residual(pts, solver.forward(pts))
-            res_loss[location] = torch.abs(target)
+            res_loss[location] = torch.abs(target).as_subclass(torch.Tensor)
             tot_loss.append(torch.abs(target))
 
         return torch.vstack(tot_loss), res_loss
@@ -74,6 +74,7 @@ def _r3_routine(self, trainer):
         """
         # compute residual (all device possible)
         tot_loss, res_loss = self._compute_residual(trainer)
+        tot_loss = tot_loss.as_subclass(torch.Tensor)
 
         # !!!!!! From now everything is performed on CPU !!!!!!
 

diff --git a/pina/label_tensor.py b/pina/label_tensor.py
@@ -118,7 +118,7 @@ def clone(self, *args, **kwargs):
         :return: a copy of the tensor
         :rtype: LabelTensor
         """
-        # used before merging 
+        # # used before merging 
         # try:
         #     out = LabelTensor(super().clone(*args, **kwargs), self.labels)
         # except:
@@ -229,8 +229,12 @@ def __getitem__(self, index):
         """
         Return a copy of the selected tensor.
         """
-        selected_lt = super(Tensor, self).__getitem__(index)
 
+        if isinstance(index, str) or (isinstance(index, (tuple, list))and all(isinstance(a, str) for a in index)):
+            return self.extract(index)
+
+        selected_lt = super(Tensor, self).__getitem__(index)
+
         try:
             len_index = len(index)
         except TypeError: