From 2a594ea2dbea5e9701c1c3494524a567f3071f83 Mon Sep 17 00:00:00 2001 From: Aryan Date: Thu, 26 Dec 2024 02:54:08 +0530 Subject: [PATCH] fix validation bug (#149) --- finetrainers/trainer.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) 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()