class_docstr: Test flow to test @pytorch_parallel.
class_docstr_tok: ["Test","flow","to","test","@","pytorch_parallel","."]
docstr: Run a simple torch parallel program where each node creates a 3 x 3 tensor
  with each entry equaling their rank + 1. Then, all reduce is called to sum the
  tensors up.
docstr_tok: ["Run","a","simple","torch","parallel","program","where","each","node","creates","a","3","x","3","tensor","with","each","entry","equaling","their","rank","+","1",".","Then",",","all","reduce","is","called","to","sum","the","tensors","up","."]
code: def parallel_step(self):
      """
      Run a simple torch parallel program where each node creates a 3 x 3 tensor
      with each entry equaling their rank + 1. Then, all reduce is called to sum the
      tensors up.
      """
      import torch
      import torch.distributed as dist
  
      # Run very simple parallel pytorch program
      dist.init_process_group(
          "gloo",
          rank=current.parallel.node_index,
          world_size=current.parallel.num_nodes,
      )
  
      # Each node creates a 3x3 matrix with values corresponding to their rank + 1
      my_tensor = torch.ones(3, 3) * (dist.get_rank() + 1)
      assert int(my_tensor[0, 0]) == current.parallel.node_index + 1
  
      # Then sum the tensors up
      print("Reducing tensor", my_tensor)
      dist.all_reduce(my_tensor, op=dist.ReduceOp.SUM)
      print("Result:", my_tensor)
  
      # Assert the values are as expected
      for i in range(3):
          for j in range(3):
              assert int(my_tensor[i, j]) == sum(
                  range(1, current.parallel.num_nodes + 1)
              )
      dist.destroy_process_group()
  
      self.node_index = current.parallel.node_index
      self.num_nodes = current.parallel.num_nodes
      self.reduced_tensor_value = int(my_tensor[0, 0])
  
      self.next(self.multinode_end)
  
imports: from metaflow import FlowSpec, step, batch, current, pytorch_parallel, Parameter
label_desc: Use image node_id 2 for calling the PytorchParallelTest obj's underlying member method code with example usage: obj.parallel_step() without return types