code: def test_lazy_apply(self) -> None:
      count_original: int = 0
      count_increment: int = 1
  
      class TestModule(LazyModuleExtensionMixin, torch.nn.Module):
          def __init__(self):
              super().__init__()
              self.count = torch.tensor(count_original)
  
          def initialize_parameters(self, input) -> None:
              pass
  
          def forward(self, input: torch.Tensor) -> torch.Tensor:
              return self.count.clone()
  
      def increment_count(module: torch.nn.Module) -> None:
          if isinstance(module, TestModule):
              module.count += torch.tensor(count_increment)
  
      def check_result(
          m: torch.nn.Module, count_after_first_forward: int
      ) -> None:
          # This check ensures that `lazy_apply()` is a delayed operation (i.e. the function is not applied immediately).
          for count in m.parameters():
              self.assertTrue(torch.allclose(count, torch.tensor(0)))
  
          input = torch.tensor(321)
          out = m(input)
  
          # This check ensures that the lazy-applied function is not called before forward function is called.
          self.assertTrue(
              torch.allclose(out, torch.tensor(count_original))
          )
  
          # This check ensures that the lazy-applied function is called after forward function is called.
          for count in m.parameters():
              self.assertTrue(
                  torch.allclose(
                      count, torch.tensor(count_after_first_forward)
                  )
              )
  
          # This check ensures that the lazy-applied function is removed after first forward pass is run.
          out = m(input)
          self.assertTrue(
              torch.allclose(
                  out, torch.tensor(count_after_first_forward)
              ),
              str(out),
          )
          # Since `increment_count` is not run the second time, value of `count` parameter is not changed.
          for count in m.parameters():
              self.assertTrue(
                  torch.allclose(
                      count, torch.tensor(count_after_first_forward)
                  )
              )
  
          # fmt: off
          check_result(
              lazy_apply(
                  TestModule(),
                  increment_count,
              ),
              count_after_first_forward=1,
          )
          check_result(
              lazy_apply(
                  torch.nn.Sequential(
                      TestModule(),
                      TestModule(),
                  ),
                  increment_count,
              ),
              count_after_first_forward=1,
          )
          check_result(
              lazy_apply(
                  lazy_apply(
                      TestModule(),
                      increment_count,
                  ),
                  increment_count,
              ),
              count_after_first_forward=2,
          )
          check_result(
              lazy_apply(
                  lazy_apply(
                      torch.nn.Sequential(
                          TestModule(),
                          TestModule()
                      ),
                      increment_count,
                  ),
                  increment_count,
              ),
              count_after_first_forward=2,
          )
          check_result(
              lazy_apply(
                  torch.nn.Sequential(
                      lazy_apply(
                          TestModule(),
                          increment_count,
                      ),
                      torch.nn.Identity()
                  ),
                  increment_count,
              ),
              count_after_first_forward=2,
          )
  
file_path: torchrec/torchrec/modules/tests/test_lazy_extension.py
imports: import inspect
  import re
  import unittest
  from typing import Tuple
  import torch
  from torch.nn.modules.lazy import LazyModuleMixin
  from torchrec.modules.lazy_extension import lazy_apply, LazyModuleExtensionMixin
label_desc: Use image node_id 5 for calling the TestLazyModuleExtensionMixin obj's underlying member method code with example usage: obj.test_lazy_apply() and returns: self