Custom function with nn.module inside nn.module class

1

Hi, using nn_module (Base class for all neural network modules. — nn_module • torch) in R I can do the following:

SomeClass <- nn_module(
     initialize = function(some_params){
     ...
     ...
     ...
     },
    forward = function(some_params){
    ...
    ... 
    ...
    },
    custom_function = function(some_params){
         nn_module(
               initialize = function(some_params){
               ...
               ...
               ...
               },
               forward = function(some_params){
               ...
               ...
               ...
              }
         )
     }
)

and so I can call custom_function() on SomeClass instances and they will return an nn_module. In PyTorch can I write inside of a custom nn.Module class a custom function which returns a nn.Module similar to what the R pseudocode above does?

I tried using a nested class and passing some_params to the __init__() function of the inner class but that did not work.

1 Like