wip

azavea · Dec 29, 2023 · c29ebe1 · c29ebe1
1 parent e902370
commit c29ebe1
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/rastervision_pytorch_learner/rastervision/pytorch_learner/learner.py b/rastervision_pytorch_learner/rastervision/pytorch_learner/learner.py
@@ -1060,7 +1060,9 @@ def build_model(self, model_def_path: Optional[str] = None) -> nn.Module:
             num_classes=cfg.data.num_classes,
             in_channels=in_channels,
             save_dir=self.modules_dir,
-            hubconf_dir=model_def_path)
+            hubconf_dir=model_def_path,
+            is_ddp_process=self.is_ddp_process,
+            is_ddp_master=self.is_ddp_master)
         return model
 
     def setup_data(self):