pytorch - 如何有条件地使用 torch.cuda.device()

Question

我有一些代码封装在:

with torch.cuda.device(self.device):
    # do a bunch of stuff

在我的 __init__ 中我有:

self.device = torch.device('cuda:0') if torch.cuda.is_available() else torch.device('cpu')

但是对于设备是cpu的情况怎么处理有点迷茫。因为 torch.cuda.device 已经明确用于 cuda。我应该只为函数写一个装饰器吗？好像有点过分了

Answer 1

根据 torch.cuda.device 的文档

device (torch.device or int) – device index to select. It’s a no-op if this argument is a negative integer or None.

基于此我们可以使用类似的东西

with torch.cuda.device(self.device if self.device.type == 'cuda' else None):
    # do a bunch of stuff

当 self.device 不是 CUDA 设备时，这只是一个空操作。