From 5de468b58320a1fadbf8bdb442204d374a76d175 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=8E=84=E9=92=9B?= Date: Sun, 18 Feb 2024 16:46:00 +0800 Subject: [PATCH] get_fs_local_rank --- olmo/torch_util.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/olmo/torch_util.py b/olmo/torch_util.py index 22afe57f5..5675e641a 100644 --- a/olmo/torch_util.py +++ b/olmo/torch_util.py @@ -55,7 +55,7 @@ def get_fs_local_rank() -> int: if all ranks share the same filesystem then `get_fs_local_rank()` will be equivalent to `get_global_rank()`, but if nodes do not share the same filesystem then `get_fs_local_rank()` will be equivalent to `get_local_rank()`. """ - return int(os.environ.get("FS_LOCAL_RANK") or get_local_rank()) + return int(os.environ.get("FS_LOCAL_RANK") or get_global_rank()) def move_to_device(o: T, device: torch.device) -> T: