ValueError: Expected input batch_size (324) to match target batch_size (4)

Could you print the shape of input before the view operation as I guess you might be changing the batch size by using view(-1, 4624). If you want to flatten the input tensor use input = input.view(input.size(0), -1) and check if you are running into shape mismatches in the first linear layer. In that case, adapt the in_features of this layer to match the activation shape.