nn.MSELoss warns you, if your output and target shapes are different, as broadcasting will be applied, which might yield wrong results. If your output has the size [batch_size, 1], make sure to pass the target is the same shape. Currently dim1 seems to be missing, so you might want to call unsqueeze(1) on the target.


The input tensor seems to be on the CPU, while a CUDATensor is expected, so you would need to push the tensor to the device.


PS: It’s better to post code snippets by wrapping them into three backticks ```, as it makes debugging easier.