From 37f021ff1be50048e7af754df303de2d2ebf4fdd Mon Sep 17 00:00:00 2001
From: Omar El Nahhas <76590695+Avic3nna@users.noreply.github.com>
Date: Tue, 26 Mar 2024 11:12:08 +0100
Subject: [PATCH] Fix multi-class heatmap plotting overview

---
 stamp/heatmaps/__main__.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/stamp/heatmaps/__main__.py b/stamp/heatmaps/__main__.py
index 1a6c95a..3c2abdc 100755
--- a/stamp/heatmaps/__main__.py
+++ b/stamp/heatmaps/__main__.py
@@ -168,7 +168,7 @@ def main(
             learn.model(feats.unsqueeze(-2), torch.ones((len(feats)))), dim=1
         )
         scores_2d = vals_to_im(scores, coords // stride).detach()
-        fig, axs = plt.subplots(nrows=2, ncols=min(2, len(categories)), figsize=(12, 8))
+        fig, axs = plt.subplots(nrows=2, ncols=max(2, len(categories)), figsize=(12, 8))
 
         show_class_map(
             class_ax=axs[0, 1],