code: def test_tuple_wait(cuda_sleep):
      # 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(
          model, [1, 1], devices=[0, 1], chunks=32, checkpoint="never"
      )
  
      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))
      y.norm().backward()
  
      torch.cuda.synchronize(0)
      torch.cuda.synchronize(1)
  
      assert torch.isclose(b.grad.norm().cpu(), torch.tensor(5.000))
  
file_path: fairscale/tests/nn/pipe/test_bugs.py
imports: import pytest
  import torch
  from torch import nn
  import torch.nn.functional
  from fairscale.fair_dev.testing.testing import skip_if_single_gpu
  from fairscale.nn.pipe import Pipe
label_desc: Use image node_id 3 for calling a global function with example usage: test_tuple_wait(cuda_sleep) and returns: x, grad, unknown,  unknown,  unknown, unknown