enable ddp find_unused_parameterse

VMarsocci · Oct 27, 2024 · 644f015 · 644f015
1 parent 3fdd72b
commit 644f015
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/pangaea/run.py b/pangaea/run.py
@@ -133,7 +133,7 @@ def main(cfg: DictConfig) -> None:
     )
     decoder.to(device)
     decoder = torch.nn.parallel.DistributedDataParallel(
-        decoder, device_ids=[local_rank], output_device=local_rank
+        decoder, device_ids=[local_rank], output_device=local_rank, find_unused_parameters=True
     )
     logger.info(
         "Built {} for with {} encoder.".format(