class_docstr: Standard solver for AudioCraft.
  
  The standard solver implements a base training loop with the following stages:
  train, valid, evaluate and generate that are expected to be all defined for
  solvers in AudioCraft. It also provides a nice default management of Dora history replay,
  checkpoint management across epoch, and logging configuration.
  
  AudioCraft solvers must inherit from the StandardSolver and define the methods
  associated to each stage as well as the show, build_model and build_dataloaders methods.
class_docstr_tok: ["Standard","solver","for","AudioCraft",".","The","standard","solver","implements","a","base","training","loop","with","the","following","stages",":","train",",","valid",",","evaluate","and","generate","that","are","expected","to","be","all","defined","for","solvers","in","AudioCraft",".","It","also","provides","a","nice","default","management","of","Dora","history","replay",",","checkpoint","management","across","epoch",",","and","logging","configuration",".","AudioCraft","solvers","must","inherit","from","the","StandardSolver","and","define","the","methods","associated","to","each","stage","as","well","as","the","show",",","build_model","and","build_dataloaders","methods","."]
docstr: Update latest best state based on pending metrics of a given stage. This method relies
  on the `BestStateDictManager.update` method to update the best state_dict with latest weights
  if the registered states happen to match to the best performing setup.
docstr_tok: ["Update","latest","best","state","based","on","pending","metrics","of","a","given","stage",".","This","method","relies","on","the","`","BestStateDictManager.update","`","method","to","update","the","best","state_dict","with","latest","weights","if","the","registered","states","happen","to","match","to","the","best","performing","setup","."]
code: def update_best_state_from_stage(self, stage_name: str = "valid"):
      """Update latest best state based on pending metrics of a given stage. This method relies
      on the `BestStateDictManager.update` method to update the best state_dict with latest weights
      if the registered states happen to match to the best performing setup.
      """
      if self.best_metric_name is None:
          # when no best metric is defined, the last state is always the best
          self._new_best_state = True
          self.logger.info("Updating best state with current state.")
      else:
          assert (
              stage_name in self._pending_metrics
          ), f"Metrics for stage {stage_name} not found."
          assert (
              self.best_metric_name in self._pending_metrics[stage_name]
          ), f"Best metric not found in {stage_name} metrics. Cannot register best state"
          current_score = self._pending_metrics[stage_name][
              self.best_metric_name
          ]
          all_best_metric_scores = [
              past_metrics[stage_name][self.best_metric_name]
              for past_metrics in self.history
          ]
          all_best_metric_scores.append(current_score)
          best_score = min(all_best_metric_scores)
          self._new_best_state = current_score == best_score
          if self._new_best_state:
              old_best = min(
                  all_best_metric_scores[:-1] + [float("inf")]
              )
              self.logger.info(
                  f"New best state with {self.best_metric_name}={current_score:.3f} (was {old_best:.3f})"
              )
  
      if self._new_best_state:
          if self.cfg.fsdp.use:
              # this will give an empty state dict on all ranks but the rank 0
              # which will have a copy in memory of the full model.
              with fsdp.switch_to_full_state_dict(self._fsdp_modules):
                  for name in self.best_state.states.keys():
                      state_source = self._get_state_source(name)
                      self.best_state.update(name, state_source)
                  # we save to a different dict.
                  self.fsdp_best_state.update(
                      self.best_state.state_dict()
                  )
              # We cannot efficiently load fsdp_best_state when using FSDP,
              # so we have do do a second pass, with the local shards.
          for name in self.best_state.states.keys():
              state_source = self._get_state_source(name)
              self.best_state.update(name, state_source)
  
imports: from abc import ABC, abstractmethod
  from contextlib import contextmanager
  from pathlib import Path
  import typing
  import flashy
  import omegaconf
  import torch
  from torch import nn
  from ..None import optim
  from ..optim import fsdp
  from ..utils import checkpoint
  from ..utils.autocast import TorchAutocast
  from ..utils.best_state import BestStateDictManager
  from ..utils.deadlock import DeadlockDetect
  from ..utils.profiler import Profiler
  from ..utils.utils import copy_state, dict_from_config, model_hash, with_rank_rng
label_desc: Use image node_id 8 for calling the StandardSolver obj's underlying member method code with example usage: obj.update_best_state_from_stage(stage_name) without return types