[tsa/eval] - fix: ensure matplotlib figures are closed after saving

- Prevent a potential memory leak by explicitly closing the plot after saving the image.
JulesBelveze · Oct 27, 2024 · 4cb8c0c · 4cb8c0c
1 parent b5bc494
commit 4cb8c0c
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/tsa/eval.py b/tsa/eval.py
@@ -50,6 +50,7 @@ def evaluate(test_iter, criterion, model, config, ts):
         plt.plot(preds, linewidth=.3)
         plt.plot(targets, linewidth=.3)
         plt.savefig("{}/preds.png".format(config.general.output_dir))
+        plt.close()
 
         torch.save(targets, os.path.join(config.general.output_dir, "targets.pt"))
         torch.save(predictions, os.path.join(config.general.output_dir, "predictions.pt"))