From 8522aa0386523e0a277e26a9732b08ce18389e10 Mon Sep 17 00:00:00 2001
From: "Michael Pilosov, PhD" <consistentbayes@gmail.com>
Date: Sat, 2 Mar 2024 23:51:05 +0000
Subject: [PATCH] isinstance

---
 scripts/sortcolor.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/scripts/sortcolor.py b/scripts/sortcolor.py
index 5d2ebfa..6d48c76 100644
--- a/scripts/sortcolor.py
+++ b/scripts/sortcolor.py
@@ -219,6 +219,8 @@ def plot_preds(
     dpi: int = 300,
     figsize=(6, 6),
 ):
+    if isinstance(preds, torch.Tensor):
+        preds = preds.detach().cpu().numpy()
     sorted_inds = np.argsort(preds.ravel())
     colors = rgb_values[sorted_inds, :3]
     if roll: