执行代码看显卡是否点亮
#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())