lint

Author: fegin

torchtitan/components/checkpoint.py

@@ -52,7 +52,7 @@ class AsyncMode(str, enum.Enum):
 # For now, we will manually pop the freqs_cis buffer, as we made this permanent
 # temporarily and we don't want to include it in the exported state_dict.
 # Context: https://github.com/pytorch/torchtitan/blob/main/torchtitan/models/llama3/model.py#L404
-excluded_parameters_for_model_only = {"freqs_cis"}
+excluded_parameters_for_model_only = {}
 
 
 class ModelWrapper(Stateful):
@@ -633,7 +633,8 @@ def _save_last_step(self, curr_step: int) -> None:
             logger.info(f"Saving a full checkpoint at last step, step {curr_step}.")
 
         save_with_gc(
-            self._flattend_model_states_sd(), checkpoint_id=self._create_checkpoint_id(curr_step)
+            self._flattend_model_states_sd(),
+            checkpoint_id=self._create_checkpoint_id(curr_step),
         )
 
     def _should_save(self, curr_step: int, last_step: bool = False) -> bool: