PyTorch - a functional equivalent of nn.Module
As we know we can wrap arbitrary number of stateful building blocks into a class which inherits from nn.Module. But how is it supposed to be done when you want to wrap a bunch of stateless functions (from nn.Functional), in order to fully utilize things which nn.Module allows you to, like automatic moving of tensors between CPU and GPU with just model.to(device)?

I already found the solution: if you have an operation inside of a module which creates a new tensor, then you have to use self.register_buffer in order to fully utilize automating moving between devices.
