fix validation bug (#149)

a-r-r-o-w · Dec 25, 2024 · 2a594ea · 2a594ea
1 parent cf9be17
commit 2a594ea
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/finetrainers/trainer.py b/finetrainers/trainer.py
@@ -813,10 +813,6 @@ def validate(self, step: int) -> None:
                 generator=self.state.generator,
             )
 
-            # Remove all hooks that might have been added during pipeline initialization to the models
-            pipeline.remove_all_hooks()
-            del pipeline
-
             prompt_filename = string_to_filename(prompt)[:25]
             artifacts = {
                 "image": {"type": "image", "value": image},
@@ -857,6 +853,10 @@ def validate(self, step: int) -> None:
                 if tracker.name == "wandb":
                     tracker.log({"validation": all_artifacts}, step=step)
 
+        # Remove all hooks that might have been added during pipeline initialization to the models
+        pipeline.remove_all_hooks()
+        del pipeline
+
         accelerator.wait_for_everyone()
 
         free_memory()