PyTorch-功能等同于nn.Module

Question

我们知道我们可以将任意数量的有状态构建块包装到从nn.Module继承的类中。但是，当您想包装一堆无状态函数（来自nn.Functional）以充分利用nn.Module允许您执行的操作（例如在CPU之间自动移动张量）时，应该怎么做？和GPU仅带有model.to(device)？

Answer 1

我已经找到了解决方案：如果模块内部有一个操作可以创建新的张量，则必须使用self.register_buffer才能充分利用设备之间的自动移动。