code: def mpi():
      seed = 1234
      torch.manual_seed(seed)
      torch.cuda.manual_seed(seed)
  
      torch.distributed.barrier()
      tensor_size = (1024, 1024, 10)
      torch.cuda.set_device(
          torch.distributed.get_rank()
      )  # need to pin device or ucx gets unhappy
  
      if torch.distributed.get_rank() == 0:
          # t = torch.Tensor(range(10)).cuda(0)
          t = torch.rand(*tensor_size).cuda(0)
          torch.distributed.send(t, 1, tag=1234)
      else:
          t = torch.empty(*tensor_size).cuda(1)
          torch.distributed.recv(t, 0, tag=1234)
          t2 = torch.rand(*tensor_size).cuda(1)
  
          assert torch.equal(t, t2)
  
imports: from collections import OrderedDict
  from copy import deepcopy
  import os
  import time
  import pytest
  import torch
  from torch import nn
  from fairscale.fair_dev.testing.testing import get_worker_map, torch_spawn
  from fairscale.internal import torch_version
  from fairscale.nn.model_parallel.initialize import get_pipeline_parallel_group
  from fairscale.nn.pipe import AsyncPipe
  from fairscale.nn.pipe.types import LazyModule
label_desc: Use image node_id 5 for calling a global function with example usage: mpi() without return types