If you define an nn.Module
, you are usually storing some submodules, parameters, buffers or other arguments in its __init__
method and write the actual forward logic in its forward
method.
This is a convenient method as nn.Module.__call__
will register hooks etc. and call finally into the forward
method.
However, you don’t need to use this approach and could completely write your model in a functional way.
I’ve written some more information in this thread which might be interesting to you.