docstr: ### Generator to load layers
code: def load(self) -> Generator[NeoXModule, None, None]:
      """
      ### Generator to load layers
      """
      with monit.section("Layers"):
          for i, (layer, files) in enumerate(self.get_layers()):
              if files is not None:
                  layer.load_state(
                      *checkpoint.load_checkpoint_files(files)
                  )
  
              layer = self.post_load_prepare(layer)
  
              monit.progress(min(0.99, (i + 1) / self.total_layers))
              yield layer
  
imports: import copy
  import math
  from typing import Dict, Optional, Set, Callable, Any, Generator, Tuple
  import torch
  from torch import nn
  from torch.cuda.amp import autocast
  from labml import monit, logger
  from labml.logger import Text
  from labml_nn.neox import checkpoint
  from labml_nn.neox.utils.cache import get_cache
label_desc: Use image node_id 11 for calling the LayerGenerator obj's underlying member method code with example usage: obj.load() without return types