1.torch.cuda()
torch.cuda会记录当前选择的GPU,并且分配的所有CUDA张量将在上面创建。可以使用torch.cuda.device上下文管理器更改所选设备。
但是,一旦张量被分配,您可以直接对其进行操作,而不考虑所选择的设备,结果将始终放在与张量相同的设备上。
默认情况下,不支持跨GPU操作,唯一的例外是copy_()。 除非启用对等存储器访问,否则对分布不同设备上的张量任何启动操作的尝试都将会引发错误。
>>> import torch
>>> x = torch.cuda.FloatTensor()
>>> x
[torch.cuda.FloatTensor with no dimension]
>>> x = torch.cuda.FloatTensor(1)
>>> x
1.00000e-02 *
-2.1046
[torch.cuda.FloatTensor of size 1 (GPU 0)]
>>> y = torch.FloatTensor(1)
>>> y
3.4179e+22
[torch.FloatTensor of size 1]
>>> y.cuda()
3.4179e+22
[torch.cuda.FloatTensor of size 1 (GPU 0)]
>> with torch.cuda.device(1):
… a = torch.cuda.FloatT