fix: use correct transforms when building val & test sets

azavea · Jan 29, 2024 · 4309ecb · 4309ecb
1 parent 67f09a6
commit 4309ecb
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/rastervision_pytorch_learner/rastervision/pytorch_learner/learner_config.py b/rastervision_pytorch_learner/rastervision/pytorch_learner/learner_config.py
@@ -1397,8 +1397,8 @@ def _build_datasets(self,
             validation, and test datasets.
         """
         train_ds = self._build_dataset('train', train_tf, tmp_dir, **kwargs)
-        val_ds = self._build_dataset('valid', train_tf, tmp_dir, **kwargs)
-        test_ds = self._build_dataset('test', train_tf, tmp_dir, **kwargs)
+        val_ds = self._build_dataset('valid', val_tf, tmp_dir, **kwargs)
+        test_ds = self._build_dataset('test', test_tf, tmp_dir, **kwargs)
         return train_ds, val_ds, test_ds
 
     def scene_to_dataset(self,