From 8be1f85c8c99e4ac5bd54690ff9e9b08a8efcfe3 Mon Sep 17 00:00:00 2001
From: Rudolph Pienaar <rudolph.pienaar@gmail.com>
Date: Thu, 25 Apr 2024 17:57:52 -0400
Subject: [PATCH] Update params for plot_bestModelOnValidate

---
 spleenseg/plotting/plotting.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/spleenseg/plotting/plotting.py b/spleenseg/plotting/plotting.py
index 9320c7c..e87adac 100644
--- a/spleenseg/plotting/plotting.py
+++ b/spleenseg/plotting/plotting.py
@@ -1,6 +1,7 @@
 #!/usr/bin/env python
 
 from pathlib import Path
+from monai.data.meta_tensor import MetaTensor
 import torch
 import matplotlib.pyplot as plt
 import pudb
@@ -42,7 +43,10 @@ def plot_trainingMetrics(
 
 
 def plot_bestModelOnValidate(
-    input: dict[str, torch.Tensor], output: torch.Tensor, title: str, savefile: Path
+    input: dict[str, torch.Tensor | MetaTensor | int],
+    output: torch.Tensor,
+    title: str,
+    savefile: Path,
 ) -> None:
     plt.figure("check", (18, 6))
     plt.subplot(1, 3, 1)