code: def python_autograd_function(pipe_class):
      # FIXME deadlock with AsyncPipe?
      # A Python autograd function might fail with this error:
      #
      #   RuntimeError: Returning Variables sharing storage with other Variables
      #   that require grad is not supported in Python functions. Please submit a
      #   feature request if you hit this error.
      #
      # It doesn't look like an essential restriction. But it happens on the
      # current PyTorch version. To avoid it, we should detach the tensor before
      # returning by identity autograd functions, such as Wait, Fork, and Join.
  
      torch.manual_seed(0)
  
      class Identity(torch.autograd.Function):
          @staticmethod
          def forward(ctx, input):
              return input
  
          @staticmethod
          def backward(ctx, grad):
              return grad
  
      class M(nn.Module):
          def forward(self, input):
              return Identity.apply(input)
  
      model = nn.Sequential(M(), M())
      model = pipe_class(
          model,
          [1, 1],
          worker_map=get_worker_map(),
          checkpoint="always",
      ).cuda()
      model.eval()
  
      x = torch.rand(42)
      y = model(x)
      if model.group.rank() == 1:
          assert torch.allclose(x, y)
  
      torch.distributed.rpc.shutdown()
      torch.distributed.barrier()
  
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 1 for calling a global function with example usage: python_autograd_function(pipe_class) and returns: input, grad, Identity