diff --git a/finetrainers/trainer.py b/finetrainers/trainer.py
index f98393c..6f2348f 100644
--- a/finetrainers/trainer.py
+++ b/finetrainers/trainer.py
@@ -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()