pre-commit

benchmarks/transformers/prob_model_text_classification.py

@@ -400,11 +400,17 @@ def unpack_model_tar(model_ckpt_path: pathlib.Path) -> pathlib.Path:
 
     model_editor = None
     if args.enable_probit_model_editor:
-        probit_freeze_fun = lambda p, v: True if "classifier" in p else False if args.probit_last_layer_only else None
+        probit_freeze_fun = (
+            lambda p, v: True
+            if "classifier" in p
+            else False
+            if args.probit_last_layer_only
+            else None
+        )
         model_editor = ProbitModelEditor(
             freeze_fun=probit_freeze_fun,
             init_log_var=args.probit_init_log_var,
-            stop_gradient=args.probit_stop_gradient
+            stop_gradient=args.probit_stop_gradient,
         )
 
     ### TRAINING

fortuna/calib_model/base.py

@@ -176,9 +176,7 @@ def load_state(self, checkpoint_dir: Path) -> None:
             )
         self.predictive.state = CalibStateRepository(checkpoint_dir=checkpoint_dir)
 
-    def save_state(
-        self, checkpoint_dir: Path, keep_top_n_checkpoints: int = 1
-    ) -> None:
+    def save_state(self, checkpoint_dir: Path, keep_top_n_checkpoints: int = 1) -> None:
         return self.predictive.state.put(
             self.predictive.state.get(),
             checkpoint_dir=checkpoint_dir,

fortuna/calib_model/calib_model_calibrator.py

@@ -13,9 +13,9 @@
 from optax._src.base import PyTree
 
 from fortuna.calib_model.state import CalibState
-from fortuna.training.trainer import TrainerABC
 from fortuna.training.mixins.jitted import JittedMixin
 from fortuna.training.mixins.multi_device import MultiDeviceMixin
+from fortuna.training.trainer import TrainerABC
 from fortuna.typing import (
     Array,
     Batch,

fortuna/likelihood/base.py

@@ -7,6 +7,7 @@
     Tuple,
     Union,
 )
+
 from flax.core import FrozenDict
 from jax import (
     jit,