Add constant scaling to coupled hpx dataset

NVIDIA · Nov 14, 2024 · 12c13da · 12c13da
1 parent ce32ad5
commit 12c13da
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/modulus/datapipes/healpix/coupledtimeseries_dataset.py b/modulus/datapipes/healpix/coupledtimeseries_dataset.py
@@ -305,11 +305,10 @@ def __getitem__(self, item):
             np.transpose(x, axes=(0, 3, 1, 2, 4, 5)) for x in inputs_result
         ]
 
-        if "constants" in self.ds.data_vars:
+
+        if self.constants is not None:
             # Add the constants as [F, C, H, W]
-            inputs_result.append(np.swapaxes(self.ds.constants.values, 0, 1))
-            # inputs_result.append(self.ds.constants.values)
-        logger.log(5, "computed batch in %0.2f s", time.time() - compute_time)
+            inputs_result.append(self.constants)
 
         # append integrated couplings
         inputs_result.append(integrated_couplings)