From 4461cdd98f841cfa936e5d3bca80c8c5dd989cdb Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Fri, 27 Dec 2024 12:25:11 +0530 Subject: [PATCH] tracker logging --- finetrainers/trainer.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/finetrainers/trainer.py b/finetrainers/trainer.py index 78e1dc3..db98e5d 100644 --- a/finetrainers/trainer.py +++ b/finetrainers/trainer.py @@ -923,7 +923,15 @@ class Trainer: tracker_key = "final" if final_validation else "validation" for tracker in accelerator.trackers: if tracker.name == "wandb": - tracker.log({tracker_key: all_artifacts}, step=step) + image_artifacts = [artifact for artifact in all_artifacts if isinstance(artifact, wandb.Image)] + video_artifacts = [artifact for artifact in all_artifacts if isinstance(artifact, wandb.Video)] + tracker.log( + { + tracker_key: {"images": image_artifacts, "videos": video_artifacts}, + }, + step=step, + ) + # Remove all hooks that might have been added during pipeline initialization to the models pipeline.remove_all_hooks()