code: def parallel_randoms(pipe_class):
      class Dropouts(nn.Module):
          def forward(self, x):
              for _ in range(100):
                  x = F.dropout(x, p=0.001)
              return x
  
      model = nn.Sequential(Dropouts(), Dropouts())
  
      x = torch.rand(10, 10, requires_grad=True).cuda()
      x.retain_grad()
      model = pipe_class(
          model,
          [1, 1],
          input_device=torch.cuda.current_device(),
          worker_map=get_worker_map(),
          chunks=10,
          checkpoint="always",
      ).cuda()
      y = model(x)
      tensor_list = [torch.empty_like(x) for _ in range(2)]
      if model.group.rank() == 1:
          y.norm().backward()
          torch.distributed.barrier()
          tensor_list[model.group.rank()] = y
          torch.distributed.all_gather(
              tensor_list, y, group=model.group
          )
          assert (
              tensor_list[0].to(torch.bool).tolist()
              == tensor_list[1].to(torch.bool).tolist()
          )
      else:
          model.back_helper(y)
          torch.distributed.barrier()
          tensor_list[model.group.rank()] = x.grad
          torch.distributed.all_gather(
              tensor_list, x.grad, group=model.group
          )
  
file_path: fairscale/tests/nn/pipe_process/test_bugs.py
imports: import pytest
  import torch
  from torch import nn
  import torch.nn.functional
  from fairscale.fair_dev.testing.testing import get_worker_map, torch_spawn
  from fairscale.nn.pipe import AsyncPipe
label_desc: Use image node_id 4 for calling a global function with example usage: parallel_randoms(pipe_class) and returns: x