sync point for distributed training only

lucaslie · lucaslie · commit eed48cbf8890 · 2020-05-08T16:15:48.000-04:00
diff --git a/src/provable_pruning/provable_pruning/util/train.py b/src/provable_pruning/provable_pruning/util/train.py
@@ -457,7 +457,8 @@ def train_with_worker(
     )
 
     # wait for all processes to load the checkpoint
-    dist.barrier()
+    if is_distributed:
+        dist.barrier()
 
     # this may be non-zero in the case of rewinding ...
     if not found_checkpoint:

Original file line number	Diff line number	Diff line change
`@@ -457,7 +457,8 @@ def train_with_worker(`
`457`	`457`	`)`
`458`	`458`
`459`	`459`	`# wait for all processes to load the checkpoint`
`460`		`- dist.barrier()`
	`460`	`+ if is_distributed:`
	`461`	`+ dist.barrier()`
`461`	`462`
`462`	`463`	`# this may be non-zero in the case of rewinding ...`
`463`	`464`	`if not found_checkpoint:`