Dimension cannot match

[Janetalready]

Hi, I tried to run train.py but got this error.

Traceback (most recent call last):
File "/home/shuwen/projects/strokenet/train.py", line 184, in
train_agent_mnist('./model/gen.pkl', './model/mnist_agent.pkl')
File "/home/shuwen/projects/strokenet/train.py", line 113, in train_agent_mnist
loss = MSE(images, approx) + penalty * LAMBDA
RuntimeError: The size of tensor a (256) must match the size of tensor b (3) at non-singleton dimension 3

It seems that MSE(images, approx) returns a shape (batch_size, 1, 256, 256) and penalty returns a shape (batch_size, 15, 3).

When I removed the arguments here MSE = torch.nn.MSELoss(reduce=False, size_average=False).to(device) and changed it to MSE = torch.nn.MSELoss().to(device), the output is just a dot. Any idea how to fix it?