I’m working on a VAE that can take in a flexible-dimension input, so I’ll like to store the pre-adaptive pooling dimensions in a variable. This is so that I can use the dimension in the decoder to help get back to the original size. I’m struggling to find a place to create this variable correctly. This is some of my code:
class Vanilla_3DVAE(nn.Module):
def __init__(self):
    super(Vanilla_3DVAE, self).__init__()
    # TODO: Parametrize the grid dimensions and channel values
    self.nn_1 = nn.Sequential(
        < layers >
        nn.AdaptiveAvgPool3d( 1 ),  # 4 -> 1x1x1
        nn.BatchNorm3d(32))
    self.nn_2 = nn.Sequential(
                              nn.Linear(32, 16),    #  32x1x1x1
                              nn.ReLU()
                            )
    self.mu_fc = nn.Linear(16,16)
    self.logvar_fc = nn.Linear(16,16)
    
    self.decoder_nn = nn.Sequential(
                                      nn.ConvTranspose3d(16,8,4,2,0),
                                      nn.BatchNorm3d(8),
                                      nn.ConvTranspose3d(8,16,4,2,0),
                                      nn.BatchNorm3d(16),
                                      nn.ConvTranspose3d(16,8,4,2,1)
                                      )

Following the suggestion of someone who helped, I separated the adaptive pooling  layer from self.nn_1, and stored the pre-pooling shape in a variable called dim in the encode() function which forward() calls upon. My new forward() code is:
def forward(self, x):
    mu,logvar, dim = self.encode(x)
    z = self.reparametrize(mu, logvar)
    decoded = self.decode(z)
    return (decoded, z, mu, logvar)

My new issue is - how can I use dim in self.decoder_nn? The __init__ cannot change - I can only take in self as the parameter. I assume that this will cause issues in using dim in self.decoder_nn.
Something else that I’ve thought of doing was
self.nn_1 = nn.Sequential( ... )
self.pre-pool_shape = x.shape
self.nn_adpative_pool = nn.Sequential(...)
self.nn_2 = nn.Sequential(....)

self.mu & logvar_fc layers

self.decoder_nn = nn.Sequential( ... <using variable pre-pool_shape> ...)

This had issues because no variable x exists in the model. If using this method, how do I declare self.pre-pool_shape?
Thank you so much