docstr: Return the state dictionary for the loss:
  - For DDP model, we rely on the default implementation of PyTorch
  - For FSDP model, we only return a shard of the teacher model and
    each rank will do the same, limiting memory consumption
docstr_tok: ["Return","the","state","dictionary","for","the","loss",":","-","For","DDP","model",",","we","rely","on","the","default","implementation","of","PyTorch","-","For","FSDP","model",",","we","only","return","a","shard","of","the","teacher","model","and","each","rank","will","do","the","same",",","limiting","memory","consumption"]
code: def state_dict(self, destination=None, prefix="", keep_vars=False):
      """
      Return the state dictionary for the loss:
      - For DDP model, we rely on the default implementation of PyTorch
      - For FSDP model, we only return a shard of the teacher model and
        each rank will do the same, limiting memory consumption
      """
      if isinstance(self.momentum_teacher, FullyShardedDataParallel):
          return {
              "center": self.center.data,
              "teacher": self.momentum_teacher.local_state_dict(),
              "teacher_meta": self.momentum_teacher.local_metadata_dict(),
          }
      else:
          return super().state_dict(destination, prefix, keep_vars)
  
imports: import logging
  import pprint
  from typing import List
  import torch
  import torch.distributed
  import torch.nn.functional
  from classy_vision.generic.distributed_util import get_cuda_device_index, get_world_size, 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 3 for calling the DINOLoss obj's underlying member method code with example usage: obj.state_dict(destination, prefix, keep_vars) and returns: dict, super