@pchandrasekaran  @ptrblck
RuntimeError: The size of tensor a (16) must match the size of tensor b (5177) at non-singleton dimension 3
with this code
class Autoencoder(nn.Module):
    def __init__(self):
        super(Autoencoder, self).__init__()
        self.encoder = nn.Sequential(
            nn.Conv2d(1, 16, 3, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(16, 32, 3, stride=2, padding=1),
            nn.ReLU(),
            nn.Conv2d(32, 64, 3)
        )
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(64, 32, 3),
            nn.ReLU(),
            nn.ConvTranspose2d(32, 16, 3, stride=2, padding=1, output_padding=1),
            nn.ReLU(),
            nn.ConvTranspose2d(16, 1, 3, stride=2, padding=1, output_padding=1),
            
        )

    def forward(self, x):
        x = self.encoder(x)
        x = self.decoder(x)
        return x


clearing the Gradients of the model parameters
optimizer.zero_grad()

# prediction for training and validation set

output_train = model(x_train.float())

output_val = model(x_val.float())


loss_train = criterion(output_train, y_train.long())

loss_val = criterion(output_val, y_val.long())