Fix lint issues

axolotl-ai-cloud · Apr 30, 2024 · 88c0326 · 88c0326
1 parent 2e8f9eb
commit 88c0326
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/src/axolotl/train.py b/src/axolotl/train.py
@@ -133,12 +133,15 @@ def terminate_handler(_, __, model_weakref):
                 _model = model_weakref()
                 if cfg.flash_optimum and BetterTransformer:
                     _model = BetterTransformer.reverse(_model)
-                _model.save_pretrained(cfg.output_dir, safe_serialization=safe_serialization)
+                _model.save_pretrained(
+                    cfg.output_dir, safe_serialization=safe_serialization
+                )
             sys.exit(0)
 
         _model_weakref = weakref.ref(model)
         signal.signal(
-            signal.SIGINT, lambda signum, frame: terminate_handler(signum, frame, _model_weakref)
+            signal.SIGINT, 
+            lambda signum, frame: terminate_handler(signum, frame, _model_weakref)
         )
 
     badge_markdown = """[<img src="https://raw.githubusercontent.com/OpenAccess-AI-Collective/axolotl/main/image/axolotl-badge-web.png" alt="Built with Axolotl" width="200" height="32"/>](https://github.com/OpenAccess-AI-Collective/axolotl)"""