What Balaji is saying is that the “reshape” operation that you wrote in your one_hot_matrix function is incorrect, but the test cases for that function did not catch your error. You can find lots of threads about this in the search history on Discourse. Here’s one. Here’s another.