code: def tuple_wait(cuda_sleep, pipe_class):
      # In v0.0.3, Wait is applied to only the first tensor on a micro-batch.
      # Under this behavior, if checkpointing was disabled, there's a possibility
      # that gradient accumulations on other tensors are not synchronized
      # properly to the copy stream.
      class Sleep(torch.autograd.Function):
          @staticmethod
          def forward(ctx, x):
              return x.detach()
  
          @staticmethod
          def backward(ctx, grad):
              with torch.cuda.device(grad.device):
                  cuda_sleep(0.05)
              return grad
  
      class Layer1(nn.Module):
          def forward(self, pair):
              a, b = pair
              return a * 1, b * 2, b * 3
  
      class Layer2(nn.Module):
          def forward(self, triple):
              a, b, c = triple
              b = Sleep.apply(b)
              return a + b + c
  
      model = nn.Sequential(Layer1(), Layer2())
      model = pipe_class(
          model,
          [1, 1],
          worker_map=get_worker_map(),
          input_device=torch.cuda.current_device(),
          chunks=32,
          checkpoint="never",
      ).cuda()
  
      a = torch.rand(1024, 3, 32, 32, device=0, requires_grad=True)
      b = torch.rand(1024, 3, 32, 32, device=0, requires_grad=True)
  
      y = model((a, b))
      if model.group.rank() == 1:
          y.norm().backward()
      else:
          model.back_helper(y)
  
      if model.group.rank() == 0:
          assert torch.isclose(b.grad.norm().cpu(), torch.tensor(5.000))
  
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 3 for calling a global function with example usage: tuple_wait(cuda_sleep, pipe_class) and returns: x, grad, unknown,  unknown,  unknown, unknown