There is no problem in the verification code that is given to you.
You have not specified the arguments to the loss function correctly, but that is not the reason for that particular error. Please check the definition of the loss function. There is something wrong with the shape of one of the labels or logits arguments. That must mean that you did the transpose incorrectly. There should be no “reshape” operation involved. They should both have shape (2, 6) by the time the loss function is called.
Also note that new_y_train is passed as the Y argument in the test case and that value was created by the output of your previous one_hot_matrix function, so the error may be there.
Here’s another thread with a checklist of common errors on this function.