code: def test_forward_lockstep():
      timeline = []
  
      class DelayedLog(nn.Module):
          def __init__(self, j, seconds):
              super().__init__()
              self.i = 0
              self.j = j
              self.seconds = seconds
  
          def forward(self, x):
              time.sleep(self.seconds)
  
              timeline.append((self.i, self.j))
              self.i += 1
  
              return x
  
      model = nn.Sequential(
          DelayedLog(0, seconds=0), DelayedLog(1, seconds=0.1)
      )
      model = Pipe(
          model, balance=[1, 1], devices=["cpu", "cpu"], chunks=3
      )
      model(torch.rand(3, 1))
  
      # Expected timeline: (Logs are recorded at !)
      #
      # Partition #0: 0! 1!   2!
      # Partition #1:    000! 111! 222!
      #
      assert timeline == [
          (0, 0),
          (1, 0),
          (0, 1),
          (2, 0),
          (1, 1),
          (2, 1),
      ]
  
imports: from collections import OrderedDict
  from copy import deepcopy
  import time
  import pytest
  import torch
  from torch import nn
  from fairscale.nn.pipe import Pipe
label_desc: Use image node_id 35 for calling a global function with example usage: test_forward_lockstep() and returns: x