code: def custom_state_dict(self, *args, **kwargs):
      state_dict = super(SimpleFrozenModel, self).state_dict(
          *args, **kwargs
      )
      custom = OrderedDict()
      for k, v in state_dict.items():
          if "linears.0.weight" not in k:
              custom[k] = v
      return custom
  
imports: import os
  import json
  import argparse
  import torch
  from collections import OrderedDict
  from deepspeed.pipe import PipelineModule, LayerSpec
  from deepspeed.moe.layer import MoE
  from deepspeed.accelerator import get_accelerator
  import deepspeed.comm
label_desc: Use image node_id 2 for calling the SimpleFrozenModel obj's underlying member method code with example usage: obj.custom_state_dict() and returns: custom