code: def __call__(self):
      # print_env()
      print("exporting PyTorch distributed environment variables")
      dist_env = submitit.helpers.TorchDistributedEnvironment().export()
      print(f"master: {dist_env.master_addr}:{dist_env.master_port}")
      print(f"rank: {dist_env.rank}")
      print(f"world size: {dist_env.world_size}")
      print(f"local rank: {dist_env.local_rank}")
      print(f"local world size: {dist_env.local_world_size}")
      # print_env()
  
      # Using the (default) env:// initialization method
      torch.distributed.init_process_group(backend="nccl")
      assert dist_env.rank == torch.distributed.get_rank()
      assert dist_env.world_size == torch.distributed.get_world_size()
  
      # Actual task / computation
      tensor = dist_env.rank * torch.ones(1).cuda()
  
      time.sleep(120)
  
      torch.distributed.all_reduce(tensor)
      if dist_env.rank == 0:
          result = list(tensor)
          print(result)
          return result
  
file_path: submitit/docs/examples/torch_distributed.py
imports: import os
  import sys
  import time
  import torch
  import submitit
label_desc: Use image node_id 1 for calling the Task obj's underlying member method code with example usage: obj.__call__() and returns: result