fix: extract freture cannot run on pure cpu

2023-04-01 17:01:51 +08:00 · 2023-04-01 17:01:51 +08:00 · 9e59375311
parent 3489b4f6da
commit 9e59375311
1 changed files with 10 additions and 5 deletions
--- a/extract_feature_print.py
+++ b/extract_feature_print.py
@ -1,9 +1,14 @@
 import os,sys,traceback
-n_part=int(sys.argv[1])
-i_part=int(sys.argv[2])
-i_gpu=sys.argv[3]
-exp_dir=sys.argv[4]
-os.environ["CUDA_VISIBLE_DEVICES"]=str(i_gpu)
+if len(sys.argv) == 4:
+    n_part=int(sys.argv[1])
+    i_part=int(sys.argv[2])
+    exp_dir=sys.argv[3]
+else:
+    n_part=int(sys.argv[1])
+    i_part=int(sys.argv[2])
+    i_gpu=sys.argv[3]
+    exp_dir=sys.argv[4]
+    os.environ["CUDA_VISIBLE_DEVICES"]=str(i_gpu)

 import torch
 import torch.nn.functional as F