pytorch中张量用GPU计算

pytorch中的张量默认存放在CPU设备中

tensor = torch.rand(3,4)

print(f"Shape of tensor: {tensor.shape}")
print(f"Datatype of tensor: {tensor.dtype}")
print(f"Device tensor is stored on: {tensor.device}")

输出:

如果GPU可用将张量移到GPU

# We move our tensor to the GPU if available
if torch.cuda.is_available():
tensor = tensor.to('cuda')