指定GPU显卡

os.environ["CUDA_VISIBLE_DEVICES"]="0"
device=torch.device('cuda' if torch.cuda.is_available() else 'cpu')