[Fix] Fix visualization hook.

2022-10-08 11:29:18 +08:00 · 2022-10-08 11:29:18 +08:00 · ae37d7fd27
parent 23cad6a0e1
commit ae37d7fd27
1 changed files with 1 additions and 1 deletions
--- a/mmcls/engine/hooks/visualization_hook.py
+++ b/mmcls/engine/hooks/visualization_hook.py
@ -78,7 +78,7 @@ class VisualizationHook(Hook):

        for sample_id in range(first_sample_id, end_idx, self.interval):
            image = images[sample_id - start_idx]
-            image = image.permute(1, 2, 0).numpy().astype('uint8')
+            image = image.permute(1, 2, 0).cpu().numpy().astype('uint8')

            data_sample = data_samples[sample_id - start_idx]
            if 'img_path' in data_sample: