python - PyTorch - nn.Module 的功能等价物

Question

正如我们所知，我们可以将任意数量的有状态构建块包装到一个继承自nn.Module. 但是，当您想包装一堆无状态函数（来自nn.Functional）时，应该如何完成，以便充分利用nn.Module允许您这样做的东西，例如在 CPU 和 GPU 之间自动移动张量model.to(device)？

标签： pythonpytorch

Answer 1

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