From 3ef67aeb063c9e153e207c8c145a8c6e333637a2 Mon Sep 17 00:00:00 2001 From: Drew McNutt Date: Mon, 22 Jul 2024 16:00:54 -0400 Subject: [PATCH] check if distributed used before world check --- python/torch_bindings.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/torch_bindings.py b/python/torch_bindings.py index d99ef1f..fd60596 100644 --- a/python/torch_bindings.py +++ b/python/torch_bindings.py @@ -329,7 +329,7 @@ def __iter__(self): worker_id = worker_info.id if worker_info is not None else 0 n_workers = worker_info.num_workers if worker_info is not None else 1 - world_size = get_world_size() + world_size = get_world_size() if torch.distributed.is_initialized() else 1 if world_size == 1: return islice(self.generate(), worker_id, None, n_workers) process_rank = get_rank()