class_docstr: Loss used in MSN (https://arxiv.org/pdf/2204.07141.pdf)
  
  Args:
      temperature (float): temperature
      teacher_momentum (float): controls EMA of teacher
class_docstr_tok: ["Loss","used","in","MSN","(","https",":","\/\/arxiv.org\/pdf\/2204.07141.pdf",")","Args",":","temperature","(","float",")",":","temperature","teacher_momentum","(","float",")",":","controls","EMA","of","teacher"]
docstr: Restore the loss state given a checkpoint
  
  Args:
      state_dict (serialized via torch.save)
docstr_tok: ["Restore","the","loss","state","given","a","checkpoint","Args",":","state_dict","(","serialized","via","torch.save",")"]
code: def load_state_dict(self, state_dict, *args, **kwargs):
      """
      Restore the loss state given a checkpoint
  
      Args:
          state_dict (serialized via torch.save)
      """
  
      # If the teacher not not been yet allocated, store it to load it
      # once the momentum teacher is available
      if self.momentum_teacher is None:
          self.checkpoint = state_dict
          logging.info("Storing the checkpoint for later use")
          return
  
      # If the momentum teacher is already allocated, use the normal
      # pytorch checkpoint restoration (load_state_dict)
      logging.info("Restoring checkpoint")
      if isinstance(self.momentum_teacher, FullyShardedDataParallel):
          self.momentum_teacher.load_local_state_dict(
              state_dict["teacher"]
          )
          self.center.copy_(state_dict["center"])
      else:
          teacher_model = get_no_ddp_model(self.momentum_teacher)
          sd = {
              x.replace("momentum_teacher.module.", ""): state_dict[x]
              for x in state_dict
              if x.find("momentum_teacher") != -1
          }
          sd = {x.replace("momentum_teacher.", ""): sd[x] for x in sd}
          teacher_model.load_state_dict(sd)
  
file_path: vissl/vissl/losses/msn_loss.py
imports: import logging
  import pprint
  from typing import List
  import torch
  from classy_vision.generic.distributed_util import get_cuda_device_index, is_distributed_training_run
  from classy_vision.losses import ClassyLoss, register_loss
  from fairscale.nn import FullyShardedDataParallel
  from vissl.config import AttrDict
  from vissl.models.model_helpers import get_no_ddp_model
label_desc: Use image node_id 4 for calling the MSNLoss obj's underlying member method code with example usage: obj.load_state_dict(state_dict) without return types