prova 6

Author: GiovanniCanali

pina/_src/solver/multi_model_simple_solver.py

@@ -206,8 +206,13 @@ def forward(self, x):
         if self._active_model_idx is not None:
             return self.models[self._active_model_idx](x)
 
+        # Strip tensor subclasses before stacking so the ensemble dimension
+        # stays unlabeled and the outer LabelTensor wrapper relabels once.
         return torch.stack(
-            [self.models[idx](x) for idx in range(self.num_models)],
+            [
+                self.models[idx](x).as_subclass(torch.Tensor)
+                for idx in range(self.num_models)
+            ],
         )
 
     def training_step(self, batch):