code: def lazy_construction(pipe_class):
      init_count = 0
  
      class Custom(nn.Module):
          def __init__(self):
              super(Custom, self).__init__()
              nonlocal init_count
              init_count += 1
  
          def forward(self, x):
              return x
  
      model = [
          LazyModule(lambda: Custom()),
          LazyModule(lambda: Custom()),
          LazyModule(lambda: Custom()),
          LazyModule(lambda: Custom()),
      ]
  
      pipe = pipe_class(
          model, balance=[2, 2], worker_map=get_worker_map()
      )
  
      assert isinstance(pipe[0], Custom)
      assert isinstance(pipe[1], Custom)
      assert len(pipe) == 2
      assert init_count == 2
  
imports: from collections import OrderedDict
  from copy import deepcopy
  import os
  import time
  import pytest
  import torch
  from torch import nn
  from fairscale.fair_dev.testing.testing import get_worker_map, torch_spawn
  from fairscale.internal import torch_version
  from fairscale.nn.model_parallel.initialize import get_pipeline_parallel_group
  from fairscale.nn.pipe import AsyncPipe
  from fairscale.nn.pipe.types import LazyModule
label_desc: Use image node_id 35 for calling a global function with example usage: lazy_construction(pipe_class) and returns: x