code: def load_state_dict(self, state_dict: Dict[str, torch.Tensor]):
      params = self.get_trainable_params()
      for n, p in params.items():
          p.data[:] = state_dict[n].to(p.data.device)
  
      for n in state_dict.keys():
          assert n in params, n
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/neox/utils/finetune.py
imports: from typing import List, Dict
  import torch
  from torch import nn
  from labml_nn.neox.model import TransformerLayer, NeoXModule
label_desc: Use image node_id 6 for calling the FineTuner obj's underlying member method code with example usage: obj.load_state_dict(state_dict) without return types