fix: extract freture cannot run on pure cpu

This commit is contained in:
源文雨 2023-04-01 17:01:51 +08:00
parent 3489b4f6da
commit 9e59375311
1 changed files with 10 additions and 5 deletions

View File

@ -1,9 +1,14 @@
import os,sys,traceback import os,sys,traceback
n_part=int(sys.argv[1]) if len(sys.argv) == 4:
i_part=int(sys.argv[2]) n_part=int(sys.argv[1])
i_gpu=sys.argv[3] i_part=int(sys.argv[2])
exp_dir=sys.argv[4] exp_dir=sys.argv[3]
os.environ["CUDA_VISIBLE_DEVICES"]=str(i_gpu) 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
import torch.nn.functional as F import torch.nn.functional as F