python基础：CUDA调用

import torch
# 指定 CPU
cpu1 =torch.device("cpu:0")
print("CPU Device:【{}:{}】".format(cpu1.type,cpu1.index))
# 指定 GPU
gpu =torch.device("cuda:0")
print("GPU Device:【{}:{}】".format(gpu.type,gpu.index))
# 建立 Tensor
device = torch.device('cuda:0')
# 使用 tensor 在 cpu 建立再转移到 gpu
gpu_tensor1= torch.Tensor([[1,1,1],[1,2,3],[1,3,5]]).to(device)  
# 使用to()方法将cup_tensor转到GPU上

或是
DEVICE = "cuda"
DEVICE_ID = "0"
CUDA_DEVICE = f"{DEVICE}:{DEVICE_ID}" if DEVICE_ID else DEVICE
if __name__ == '__main__':
    tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True)
    model = AutoModel.from_pretrained("THUDM/chatglm2-6b", trust_remote_code=True).cuda()

验证cuda

python

import torch

#测试版本

print(torch.__version__)

print(torch.version.cuda)

print(torch.cuda.is_available())