fix validation bug (#149)

2026-10-09 00:09:45 +00:00 · 2024-12-26 02:54:08 +05:30
parent cf9be17819
commit 2a594ea2db
1 changed files with 4 additions and 4 deletions
@@ -813,10 +813,6 @@ class Trainer:
                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 @@ class Trainer:
                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()