This MSE should be 1.0 — every prediction is exactly 1 away from its target. It prints 2.333, built from NINE loss elements for three samples: the (3,) predictions were compared against the (3, 1) column of targets pairwise, the silent (3, 3) broadcast in list form. Fix it by flattening the targets to (3,) and pairing predictions with their own targets only.
The break is on line 5 — but read the whole snippet first.