Skip to content

Commit 9e59375

Browse files
committed
fix: extract freture cannot run on pure cpu
1 parent 3489b4f commit 9e59375

File tree

1 file changed

+10
-5
lines changed

1 file changed

+10
-5
lines changed

extract_feature_print.py

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,14 @@
11
import os,sys,traceback
2-
n_part=int(sys.argv[1])
3-
i_part=int(sys.argv[2])
4-
i_gpu=sys.argv[3]
5-
exp_dir=sys.argv[4]
6-
os.environ["CUDA_VISIBLE_DEVICES"]=str(i_gpu)
2+
if len(sys.argv) == 4:
3+
n_part=int(sys.argv[1])
4+
i_part=int(sys.argv[2])
5+
exp_dir=sys.argv[3]
6+
else:
7+
n_part=int(sys.argv[1])
8+
i_part=int(sys.argv[2])
9+
i_gpu=sys.argv[3]
10+
exp_dir=sys.argv[4]
11+
os.environ["CUDA_VISIBLE_DEVICES"]=str(i_gpu)
712

813
import torch
914
import torch.nn.functional as F

0 commit comments

Comments
 (0)