Could you print the shape of y_pred and y_data?
Both should have the same shape, i.e. [batch_size, 1] in your case.
In the past we’ve had similar issues where dim1 was missing in the target tensor, so that an unwanted broadcasting was performed. The latest PyTorch version will throw a warning, if this shape mismatch is detected.