Fix validation set loading during fit stage

PTG-Kitware · Nov 20, 2024 · 1fa96e8 · 1fa96e8
1 parent 7dcef0e
commit 1fa96e8
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tcn_hpl/data/ptg_datamodule.py b/tcn_hpl/data/ptg_datamodule.py
@@ -177,7 +177,7 @@ def setup(self, stage: Optional[str] = None) -> None:
                 kwcoco.CocoDataset(self.hparams.coco_train_poses),
                 self.hparams.target_framerate,
             )
-        if stage == "validate" and not self.data_val:
+        if (stage == "validate" or stage == "fit") and not self.data_val:
             self.data_val.load_data_offline(
                 kwcoco.CocoDataset(self.hparams.coco_validation_activities),
                 kwcoco.CocoDataset(self.hparams.coco_validation_objects),