submitit example

pytorch · Nov 23, 2023 · 9132a60 · 9132a60
1 parent 89770e4
commit 9132a60
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 6 deletions.
diff --git a/examples/impala/config_multi_node_submitit.yaml b/examples/impala/config_multi_node_submitit.yaml
@@ -17,7 +17,7 @@ collector:
   backend: gloo
   frames_per_batch: 80
   total_frames: 200_000_000
-  num_workers: 4
+  num_workers: 1
 
 # logger
 logger:

diff --git a/examples/impala/impala_multi_node_submitit.py b/examples/impala/impala_multi_node_submitit.py
@@ -73,7 +73,6 @@ def main(cfg: "DictConfig"):  # noqa: F821
         raise NotImplementedError(
             f"device assignment not implemented for backend {cfg.collector.backend}"
         )
-    import ipdb; ipdb.set_trace()
     collector = DistributedDataCollector(
         create_env_fn=[make_env(cfg.env.env_name, device)] * num_workers,
         policy=actor,
@@ -85,11 +84,9 @@ def main(cfg: "DictConfig"):  # noqa: F821
         slurm_kwargs=slurm_kwargs,
         storing_device="cuda:0" if cfg.collector.backend == "nccl" else "cpu",
         launcher="submitit",
-        update_after_each_batch=True,
+        # update_after_each_batch=True,
         backend=cfg.collector.backend,
-        )
-
-    import ipdb; ipdb.set_trace()
+    )
 
     # Create data buffer
     sampler = SamplerWithoutReplacement()