I have the following training data code which corresponds to a simple quantum mechanics example where I consider a Y-rotated qubit, and generate measurement outcomes, either spin-up or spin-down for 10 discretized rotations $\theta_j \in (0, \pi)$. For each rotation angle I generate a set of 100 measurements, which will either be spin-up or spin-down, I then tally these measurements as tuples which indicate the total spin-up and spin-down qubits observed for each rotation angle $\theta_j$ and use them as a training set by associating them with one-hot vectors. After training, when I test this with say the result {{97,3}} I would expect to get the first one-hot vector corresponding to the first rotation angle but it instead outputs the third as shown in the code. Can anyone advise on what is going wrong with the training in this code. Many thanks.