Ok good to know. And here is the current state of my code:
import torch
import torch.nn as nn
import torch.nn.init as init
from torch.autograd import Variable
from torchvision import models
from torchsummary import summary

def reparametrize(mu, logvar):
std = logvar.div(2).exp()
eps = Variable(std.data.new(std.size()).normal_())
return mu + std*eps

class View(nn.Module):
def __init__(self, size):
    super(View, self).__init__()
    self.size = size

def forward(self, tensor):
    return tensor.view(self.size)


class Print(nn.Module):
def __init__(self):
    super(Print, self).__init__()

def forward(self, x):
    print(x.shape)
    return x

class BetaVAE_H(nn.Module):
"""Model proposed in original beta-VAE paper(Higgins et al, ICLR, 2017)."""

def __init__(self, z_dim=10, nc=8):
    super(BetaVAE_H, self).__init__()
    self.z_dim = z_dim
    self.nc = nc
    self.encoder = nn.Sequential(
        nn.Conv2d(nc, 32, 4, 2, 1),          # B,  32, 48, 48
        Print(),
        nn.ReLU(True),
        nn.Conv2d(32, 32, 4, 2, 1),          # B,  32, 24, 24
        Print(),
        nn.ReLU(True),
        nn.Conv2d(32, 64, 4, 2, 1),          # B,  64,  12, 12
        Print(),
        nn.ReLU(True),
        nn.Conv2d(64, 64, 4, 2, 1),          # B,  64,  6,  6
        Print(),
        nn.ReLU(True),
        nn.Conv2d(64, 256, 4, 1),            # B, 256,  3,  3
        Print(),
        nn.ReLU(True),
        View((1,-1)),                 # B, 256
        Print(),
        nn.Linear(256*8*8,z_dim*2),             # B, z_dim*2
        Print(),
    )
    self.decoder = nn.Sequential(
        Print(),
        nn.Linear(z_dim, 256*3*3),               # B, 256
        Print(),
        View((1, -1, 3, 3)),               # B, 256,  1,  1
        Print(),
        nn.ReLU(True),
        nn.ConvTranspose2d(256, 64, 4),      # B,  64,  4,  4
        Print(),
        nn.ReLU(True),
        nn.ConvTranspose2d(64, 64, 4, 2, 1), # B,  64,  8,  8
        Print(),
        nn.ReLU(True),
        nn.ConvTranspose2d(64, 32, 4, 2, 1), # B,  32, 16, 16
        Print(),
        nn.ReLU(True),
        nn.ConvTranspose2d(32, 32, 4, 2, 1), # B,  32, 32, 32
        Print(),
        nn.ReLU(True),
        nn.ConvTranspose2d(32, nc, 4, 2, 1),  # B, nc, 64, 64
        Print(),
        nn.Upsample(96,mode='bilinear'),

    )

    self.weight_init()

def weight_init(self):
    for block in self._modules:
        for m in self._modules[block]:
            kaiming_init(m)

def forward(self, x):
    #print("x val: " + str(x.size()))
    distributions = self._encode(x)
    mu = distributions[:, :self.z_dim]
    logvar = distributions[:, self.z_dim:]
    z = reparametrize(mu, logvar)
    #print("z val: " + str(z.size()))
    x_recon = self._decode(z)

    return x_recon, mu, logvar

def _encode(self, x):
    return self.encoder(x)

def _decode(self, z):
    return self.decoder(z)

def kaiming_init(m):
if isinstance(m, (nn.Linear, nn.Conv2d)):
    init.kaiming_normal(m.weight)
    if m.bias is not None:
        m.bias.data.fill_(0)
elif isinstance(m, (nn.BatchNorm1d, nn.BatchNorm2d)):
    m.weight.data.fill_(1)
    if m.bias is not None:
        m.bias.data.fill_(0)