-
Notifications
You must be signed in to change notification settings - Fork 771
Description
I have a module/class of this kind:
class autoencoder(nn.Module):
hidden_dim: int
z_dim: int
output_dim: int
def setup(self):
self.encoder = encoder(self.hidden_dim, self.z_dim)
self.decoder = decoder(self.hidden_dim, self.output_dim)
def __call__(self, x):
z = self.encoder(x)
y = self.decoder(z)
return y
Here, encoder and decoder are classes defined elsewhere. On occasions, I want to be able to call the individual submodules (e.g. self.encoder or self.decoder) of the module autoencoder, but as I understand they are inaccessible and can only be accessed by calling __ call __ (e.g. via model.apply), which calls both of them, consecutively (which I do also want to do sometimes). I ideally want to keep these submodules bundled together in this module so that when I initialize the model (via model.init), I have the full model and all of it's parameters in one place. (In my actual problem I have 6 submodules so I have simplified things here.)
This use case seems to be described here under 'Future work'. Is there no current solution to what I am trying to do?