测试pytorch-gpu

执行代码看显卡是否点亮

#coding=gbk
import torch

# 定义张量的形状和大小
shape = (100, 1000)
num_tensors = 50000


device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
data = [torch.rand(shape, device=device) for _ in range(num_tensors)]

total_sum = torch.tensor([0.0])
for tensor in data:
    total_sum += tensor.sum().cpu()

print('Total sum:', total_sum.item())