diff --git a/src/callbacks_wandb.py b/src/callbacks_wandb.py
index bd2b917e..3602732f 100644
--- a/src/callbacks_wandb.py
+++ b/src/callbacks_wandb.py
@@ -235,3 +235,4 @@ def on_validation_end(
                 axs[1, i].axis("off")
 
             self.logger.experiment.log({"Images": wandb.Image(fig)})
+            plt.close(fig)
diff --git a/trainer.py b/trainer.py
index ef2f4425..92b15a9a 100644
--- a/trainer.py
+++ b/trainer.py
@@ -9,6 +9,7 @@
 - https://lightning.ai/docs/pytorch/2.1.0/cli/lightning_cli.html
 - https://pytorch-lightning.medium.com/introducing-lightningcli-v2-supercharge-your-training-c070d43c7dd6
 """
+
 from lightning.pytorch.callbacks import (
     LearningRateMonitor,  # noqa: F401
     ModelCheckpoint,
@@ -75,4 +76,5 @@ def cli_main(
 # %%
 if __name__ == "__main__":
     cli_main()
+
     print("Done!")