From e547c6261a83eb8976fd5fbc2c359a21c90baee4 Mon Sep 17 00:00:00 2001
From: MichaelAnderson-AI <michael.anthoney.anderson@gmail.com>
Date: Wed, 7 Jul 2021 16:23:31 +0200
Subject: [PATCH] `feature_visualization()` CUDA fix (#3925)

---
 utils/plots.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/utils/plots.py b/utils/plots.py
index 1ab3bb6..23a4862 100644
--- a/utils/plots.py
+++ b/utils/plots.py
@@ -461,7 +461,7 @@ def feature_visualization(x, module_type, stage, n=64, save_dir=Path('runs/detec
             f = f"stage{stage}_{module_type.split('.')[-1]}_features.png"  # filename
 
             plt.figure(tight_layout=True)
-            blocks = torch.chunk(x[0], channels, dim=0)  # select batch index 0, block by channels
+            blocks = torch.chunk(x[0].cpu(), channels, dim=0)  # select batch index 0, block by channels
             n = min(n, channels)  # number of plots
             ax = plt.subplots(math.ceil(n / 8), 8, tight_layout=True)[1].ravel()  # 8 rows x n/8 cols
             for i in range(n):