code: def on_epoch(self, where: float) -> None:
      # Run the normal LR schedulers
      super().on_epoch(where)
  
      # Materialize the optimizer state on the replica in charge of checkpointing
      logging.info(
          "Consolidating sharded state on primary rank. Where: %d"
          % where
      )
      self.consolidate_state_dict()
  
file_path: ClassyVision/classy_vision/optim/zero.py
imports: import logging
  from typing import Any
  import torch.distributed
  from classy_vision.generic.distributed_util import get_primary_rank
  from classy_vision.optim import build_optimizer, ClassyOptimizer, register_optimizer
label_desc: Use image node_id 4 for calling the ZeRO obj's underlying member method code with example usage: obj.on_epoch(where) without return types