blip2 stage1 losses

Summary: as title Reviewed By: ankitade Differential Revision: D48103991
facebookresearch · Sep 5, 2023 · c59e2d7 · c59e2d7
1 parent 1b4f79f
commit c59e2d7
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/torchmultimodal/utils/distributed.py b/torchmultimodal/utils/distributed.py
@@ -39,3 +39,10 @@ def gather_tensor(tensor: Tensor, backprop_in_gather: bool = True) -> List[Tenso
         all_gather_no_backprop(tensor_all_gpus, tensor)
         tensor_all_gpus[torch.distributed.get_rank()] = tensor
         return tensor_all_gpus
+
+
+def get_rank():
+    """get rank util for distributed training"""
+    if torch.distributed.is_available() and torch.distributed.is_initialized():
+        return torch.distributed.get_rank()
+    return 0