把張量計算放到GPU上進行計算

先設置環境

parser = argparse.ArgumentParser()
parser.add_argument('--gpu',type=str,default='7')
opt=parser.parse_args()
print(opt)
os.environ["CUDA_VISIBLE_DEVICES"]=opt.gpu



然后把要放到GPU中的變量變成cuda變量

image_torch = image_torch.cuda()
其中變量要為tensor才行

計算過程中自然產生的變量不用再.cuda（）放到GPU上，已經自動放上去了

但是新聲明的變量要放上去，比如中間過程中：

S_out = torch.zeros(S_max.size())
就需要S_out = S_out.cuda（）放到GPU上



最后把計算結果放回到CPU上

S_mean = S_mean.cpu()
S_mean = S_mean.numpy().tolist()