using tensor device in eval distributed (#557)

* using tensor device in eval distributed * tensor_to_rank0 device
Oneflow-Inc · Oct 28, 2024 · 9dcbe3b · 9dcbe3b
1 parent 13f1b12
commit 9dcbe3b
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 3 deletions.
diff --git a/libai/evaluation/evaluator.py b/libai/evaluation/evaluator.py
@@ -203,12 +203,12 @@ def inference_on_dataset(
 
             # get valid sample
             valid_data = {
-                key: dist.tensor_to_rank0(value, to_local=True)[:valid_sample]
+                key: dist.tensor_to_rank0(value, device=value.placement.type, to_local=True)[:valid_sample]
                 for key, value in data.items()
             }
             valid_outputs = {}
             for key, value in outputs.items():
-                value = dist.tensor_to_rank0(value, to_local=True)
+                value = dist.tensor_to_rank0(value, device=value.placement.type, to_local=True)
                 if value.ndim > 1:
                     valid_outputs[key] = value[:valid_sample]  # Slice if it's batched output
                 else:

diff --git a/libai/utils/distributed.py b/libai/utils/distributed.py
@@ -471,7 +471,6 @@ def tton(tensor, local_only=False, ranks=None):
 
 def tensor_to_rank0(tensor, device="cuda", to_local=False):
     """Global tensor to rank0."""
-    # assert device in ["cpu", "cuda"], f"not supported for device:{device}"
     if tensor.is_global:
         # Consider if it's 2d mesh, ranks should be [[0]] instead of [0]
         placement = flow.placement(device, ranks=[0] if tensor.placement.ranks.ndim == 1 else [[0]])