Fix multi-class heatmap plotting overview

KatherLab · Mar 26, 2024 · 37f021f · 37f021f
1 parent 2f0a143
commit 37f021f
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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],