self.features would output a tensor of [batch_size, 512, 14, 14].
If you reshape this output via view(-1, 512, 12*14*14), you are interleaving the data.
What’s the use case to put the batch dimension into the features?
This also means that you end up with a batch size of 1, which creates the mentioned error.