diff --git a/finetrainers/trainer.py b/finetrainers/trainer.py
index 3179586..ad03179 100644
--- a/finetrainers/trainer.py
+++ b/finetrainers/trainer.py
@@ -884,10 +884,6 @@ class Trainer:
                 # todo support passing `fps` for supported pipelines.
             )
 
-            # 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},
@@ -930,6 +926,10 @@ class Trainer:
                 if tracker.name == "wandb":
                     tracker.log({tracker_key: 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()