I am writing a flexible VAE that takes in inputs of varying sizes. This is the code that I have so far:
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(
            nn.Conv3d(8, 16, 4, 2, 1 ),        # -4 from dimensions
            nn.ReLU(),
            nn.AvgPool3d(2),            # /2
            nn.BatchNorm3d(16),
            nn.Conv3d(16, 32, 4, 2, 1 ),       # -4
            nn.ReLU(),
            nn.AdaptiveAvgPool3d( 1 ),  # dimension -> 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)
                                          )

Since the inputs will have different dimensions, I’d like to keep track of the dimension size right before nn.AdaptiveAvgPool3d( 1 ), storing this in a variable. I don’t have access to the input dimension size in this .py file. What is the best way to do this?