Update test_heatmap.py

cz1236 · May 11, 2024 · a259603 · a259603
1 parent 5976590
commit a259603
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/FTransUNet/test_heatmap.py b/FTransUNet/test_heatmap.py
@@ -29,6 +29,7 @@
 handle = nvmlDeviceGetHandleByIndex(int(os.environ["CUDA_VISIBLE_DEVICES"]))
 print("Device :", nvmlDeviceGetName(handle))
 
+BATCH_SIZE = 1
 config_vit = CONFIGS_ViT_seg['R50-ViT-B_16']
 config_vit.n_classes = 6
 config_vit.n_skip = 3