Hello. I’m studying neural networks, and the task is to create a NN with this architecture. But the thing is that the error comes in

ValueError: Expected input batch_size (300) to match target batch_size (100).

According to assignment, I need to use MSELoss. Tried to fix it by transposing output values, but nothing works. Tried to change function to CrossEntropyLoss, but no luck either. Transposing helped, but the accuracy values did not increase and the error did not drop.
What could be the problem?
How do I get the correct architecture in such a trivial problem without transpositions and get an improvement?
It’s worth noting that the NN works on MNIST without any problems. This only happens on CIFAR10.
Link to Google Colab: https://colab.research.google.com/drive/1IIVlUoUgcpHYv3psChjLlmqyZu_XI98q?usp=sharing 3