Add description for reload func

neuralmagic · Jun 30, 2023 · d19c550 · d19c550
1 parent 80620f1
commit d19c550
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/sparseml/transformers/sparsification/trainer.py b/src/sparseml/transformers/sparsification/trainer.py
@@ -653,6 +653,10 @@ def _setup_manager(
         return manager, arch_manager
 
     def _reload_model_state(self, load_path: str, orig_state_dict: Dict[str, Any]):
+        """
+        Reload the weights after model arch changes due to recipe application
+        Return True if weights are successfully reloaded; False otherwise
+        """
         invalid_load_path = not load_path or not os.path.isdir(load_path)
         files = os.listdir(load_path) if not invalid_load_path else []
         weight_files = [