[fix]

main.py

@@ -12,8 +12,8 @@ import torch.backends.cudnn as cudnn
 import re
 from utils import *
 
-# os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3,4"
-os.environ["CUDA_VISIBLE_DEVICES"] = "0"
+os.environ["CUDA_VISIBLE_DEVICES"] = "0,1,2,3"
+# os.environ["CUDA_VISIBLE_DEVICES"] = "0"
 
 
 def parse_option():
@@ -173,7 +173,7 @@ def main():
     elif args.dataset == "food2k":
         NUM_CATEGORIES = 2000
     elif args.dataset == "jkyy":
-        NUM_CATEGORIES = 1788
+        NUM_CATEGORIES = 1790
     elif args.dataset == "test":
         NUM_CATEGORIES = 5
 
@@ -214,8 +214,8 @@ def main():
 
     cudnn.benchmark = True
     net.cuda()
-    # device_ids = [0, 1, 2, 3, 4]
-    device_ids = [0]
+    device_ids = [0, 1, 2, 3]
+    # device_ids = [0]
     # net = nn.DataParallel(net).to(device_ids)
     net = nn.DataParallel(net, device_ids=device_ids)
     # optimizer = nn.DataParallel(optimizer, device_ids=device_ids)