[fix] Fix for dlwp_healpix dataloader not subsetting channels correct…

…ly (#683) * Fix for dlwp_healpix dataloader not subsetting channels correctly * Adding fix for coupled dataloader as well * format fix
NVIDIA · Oct 29, 2024 · f46e25f · f46e25f
1 parent f7a5773
commit f46e25f
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/modulus/datapipes/healpix/data_modules.py b/modulus/datapipes/healpix/data_modules.py
@@ -568,6 +568,10 @@ def setup(self) -> None:
                     batch_size=self.batch_size,
                 )
 
+                dataset = dataset.sel(
+                    channel_in=self.input_variables,
+                    channel_out=self.output_variables,
+                )
             else:
                 dataset = open_fn(
                     input_variables=self.input_variables,
@@ -984,6 +988,10 @@ def setup(self) -> None:
                     batch_size=self.batch_size,
                 )
 
+                dataset = dataset.sel(
+                    channel_in=self.input_variables,
+                    channel_out=self.output_variables,
+                )
             else:
                 dataset = open_fn(
                     input_variables=self.input_variables + coupled_variables,