Improve logic of final plotter

FNNDSC · Apr 25, 2024 · 18deddf · 18deddf
1 parent 1c5a101
commit 18deddf
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/spleenseg/plotting/plotting.py b/spleenseg/plotting/plotting.py
@@ -71,6 +71,6 @@ def plot_infer(
     plt.imshow(input[:, :, 20], cmap="gray")
     plt.subplot(1, 2, 2)
     plt.title(f"result: {title}")
-    plt.imshow(result[0].detach().cpu()[1, :, :, 20])
+    plt.imshow(result.detach().cpu()[1, :, :, 20])
     plt.savefig(str(savefile))
     plt.clf()