From 07268f02bd228829bf1eb987b6839954f17f23d9 Mon Sep 17 00:00:00 2001
From: Tim Hoffmann <2836374+timhoffm@users.noreply.github.com>
Date: Sun, 25 Sep 2022 22:34:17 +0200
Subject: [PATCH] Fix evaluating colormaps on non-numpy arrays

Closes #23132.

I can't specifically test that case, because we don't have pytorch as a
test dependency. However, I'd claim that deferring np.nan() to after
converting to a numpy array is obviously the right thing to do.
---
 lib/matplotlib/colors.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/lib/matplotlib/colors.py b/lib/matplotlib/colors.py
index 5441b0d617b5..f8a18443c9e0 100644
--- a/lib/matplotlib/colors.py
+++ b/lib/matplotlib/colors.py
@@ -706,8 +706,12 @@ def __call__(self, X, alpha=None, bytes=False):
         if not self._isinit:
             self._init()
 
-        mask_bad = X.mask if np.ma.is_masked(X) else np.isnan(X)  # Mask nan's.
+        # Take the bad mask from a masked array, or in all other cases defer
+        # np.isnan() to after we have converted to an array.
+        mask_bad = X.mask if np.ma.is_masked(X) else None
         xa = np.array(X, copy=True)
+        if mask_bad is None:
+            mask_bad = np.isnan(xa)
         if not xa.dtype.isnative:
             xa = xa.byteswap().newbyteorder()  # Native byteorder is faster.
         if xa.dtype.kind == "f":