Thanks for replying :). Yes I’d like to somehow pass the dimension back into  self.decoder_nn = nn.Sequential( ... ) and use the dimension to help get to the original size. The code that I had to get the dimension in in encode():
class Vanilla_3DVAE(nn.Module):
    def __init__(self):
        super(Vanilla_3DVAE, self).__init__()
        self.nn_1 = nn.Sequential( ... ) # same as before except missing last 2 layers, moved below
        self.nn_adaptivePool = nn.Sequential(
            nn.AdaptiveAvgPool3d( 1 ),  # 4 -> 1x1x1
            nn.BatchNorm3d(32))
        self.nn_2 = nn.Sequential( ... ) # same as before
        self.mu_fc = nn.Linear(16,16)
        self.logvar_fc = nn.Linear(16,16)
        
        self.decoder_nn = nn.Sequential( ... ) # same as before but would like to pass the dimension here

    def encode(self, x):
        encoded = self.nn_1(x)
        prePoolShape = list( encoded.size())
        prePoolDim = prePoolShape[ 2 ] # <<<<< dimension <<<<<<<<<
        # print( prePoolDim ) 
        encoded1= self.nn_adaptivePool( encoded )
        encoded_flat = self.flatten(encoded1 )
        encoded_flat = self.nn_2(encoded_flat)
        mu = self.mu_fc(encoded_flat)
        logvar = self.logvar_fc(encoded_flat)
        return mu, logvar

    def decode(self, z ):
        z_4d = self.restore_dim(z)
        return self.decoder_nn(z_4d )

    def forward(self, x):
        mu,logvar = self.encode(x)
        z = self.reparametrize(mu, logvar)
        decoded = self.decode(z)
        return (decoded, z, mu, logvar)