小编hsh*_*hsh的帖子

nn.Module.cuda() 将所有模型参数和缓冲区移动到 GPU。

但是为什么不是模型成员张量呢？

class ToyModule(torch.nn.Module):
    def __init__(self) -> None:
        super(ToyModule, self).__init__()
        self.layer = torch.nn.Linear(2, 2)
        self.expected_moved_cuda_tensor = torch.tensor([0, 2, 3])

    def forward(self, input: torch.Tensor) -> torch.Tensor:
        return self.layer(input)

toy_module = ToyModule()
toy_module.cuda()

next(toy_module.layer.parameters()).device
>>> device(type='cuda', index=0)

对于模型成员张量，设备保持不变。

>>> toy_module.expected_moved_cuda_tensor.device
device(type='cpu')