My vae class looks like this:
class Encoder(nn.Module): def __init__(self): super(Encoder, self).__init__() c = capacity self.conv1 = nn.Conv2d(in_cha