PyTorch - функциональный эквивалент nn.Module

Question

Как мы знаем, мы можем заключить произвольное количество блоков с сохранением состояния в класс, который наследуется от nn.Module. Но как это должно быть сделано, когда вы хотите обернуть кучу функций без состояния (от nn.Functional), чтобы полностью использовать вещи, которые nn.Module позволяет вам, например, автоматическое перемещение тензоров между CPU и GPU с помощью всего лишьmodel.to(device)

Answer 1

Я уже нашел решение: если у вас есть операция внутри модуля, которая создает новый тензор, то вы должны использовать self.register_buffer, чтобы полностью использовать автоматизацию перемещения между устройствами.