def consolidate_optimizer(self):
        """For OSS, we need to consolidate the state dict."""
        if self.cfg.checkpoint.no_save_optimizer_state:
            return
        self._gathered_optim_state = None
        if hasattr(self.optimizer.optimizer, "consolidate_state_dict"):
            self.optimizer.optimizer.consolidate_state_dict()
        elif self.is_fsdp and not self.model.use_sharded_state:
            st = self.model.gather_full_optim_state_dict(
                self.optimizer
            )  # only returns on rank 0
            self._gathered_optim_state = st
For OSS, we need to consolidate the state dict.
structured_prompting/fairseq-version/fairseq/fairseq/trainer.py