fix typo

ewencedr · Jun 14, 2023 · 6318a73 · 6318a73
1 parent 46bcfc0
commit 6318a73
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/src/data/jetnet_datamodule.py b/src/data/jetnet_datamodule.py
@@ -287,9 +287,11 @@ def setup(self, stage: Optional[str] = None):
             mask_test = torch.tensor(np.expand_dims(mask_test[..., 0], axis=-1))
 
             if self.hparams.normalize:
-                self.data_train = TensorDataset(tensor_train, mask_train)
-                self.data_val = TensorDataset(tensor_val, mask_val)
-                self.data_test = TensorDataset(tensor_test, mask_test)
+                self.data_train = TensorDataset(
+                    tensor_train, mask_train, tensor_conditioning_train
+                )
+                self.data_val = TensorDataset(tensor_val, mask_val, tensor_conditioning_val)
+                self.data_test = TensorDataset(tensor_test, mask_test, tensor_conditioning_test)
 
                 self.means = torch.tensor(means)
                 self.stds = torch.tensor(stds)
@@ -315,7 +317,7 @@ def setup(self, stage: Optional[str] = None):
             self.mask_val = unnormalized_mask_val
             self.x_mean = x_mean
             self.x_cov = x_cov
-            self.tensor_conditioning_train = tensor_conditioning_test
+            self.tensor_conditioning_train = tensor_conditioning_train
             self.tensor_conditioning_val = tensor_conditioning_val
             self.tensor_conditioning_test = tensor_conditioning_test