class_docstr: Test the PipelineHelper class.
class_docstr_tok: ["Test","the","PipelineHelper","class","."]
code: def test_schedule_work_items(self):
      # test that we schedule things correctly
      # pretend we have 8 layers and 4 gpus, and they are unevenly distributed
      model = torch.nn.ModuleList()
      for i in range(8):
          layer = IdentityLayer()
          if i == 0:
              layer._mp_gpu = "cuda:0"
          elif i in (1, 2, 3):
              layer._mp_gpu = "cuda:1"
          elif i in (4, 5):
              layer._mp_gpu = "cuda:2"
          elif i in (6, 7):
              layer._mp_gpu = "cuda:3"
          model.append(layer)
  
      # there are 2 chunks, each 16 x 7 in size
      chunks = PipelineHelper.split(torch.randn(32, 7), 16)
  
      work_items = list(
          PipelineHelper.schedule_work_items(model, chunks)
      )
      assert len(work_items) == 8
      assert (
          work_items[0].layer_nos == [0]
          and work_items[0].chunk_idx == 0
      )
      assert (
          work_items[1].layer_nos == [1, 2, 3]
          and work_items[1].chunk_idx == 0
      )
      assert (
          work_items[2].layer_nos == [0]
          and work_items[2].chunk_idx == 1
      )
      assert (
          work_items[3].layer_nos == [4, 5]
          and work_items[3].chunk_idx == 0
      )
      assert (
          work_items[4].layer_nos == [1, 2, 3]
          and work_items[4].chunk_idx == 1
      )
      assert (
          work_items[5].layer_nos == [6, 7]
          and work_items[5].chunk_idx == 0
      )
      assert (
          work_items[6].layer_nos == [4, 5]
          and work_items[6].chunk_idx == 1
      )
      assert (
          work_items[7].layer_nos == [6, 7]
          and work_items[7].chunk_idx == 1
      )
  
imports: import torch
  import numpy
  import unittest
  from parlai.utils.torch import padded_tensor, argsort, PipelineHelper, neginf, IdentityLayer, trainable_parameters, total_parameters
label_desc: Use image node_id 12 for calling the TestPipelineHelper obj's underlying member method code with example usage: obj.test_schedule_work_items() without return types