torch显卡操作

import torch
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0,1"

device0 = torch.device('cuda:0')
device1 = torch.device('cuda:1')
data = torch.rand(300,300)

print(data)
print(torch.cuda.device_count()) # 显卡数量 2
print(torch.cuda.get_device_name()) # 设备名 NVIDIA GeForce RTX 3090
print(torch.cuda.current_device()) # 当前设备编号 0
data0 = data.to(device0)
data1 = data.to(device1)
print(data0.device) # cuda:0
print(data1.device) # cuda:1
time.sleep(60)