How to encode positional info in a 2D matrix input to PyTorch

I’m working on a project where the goal is to pass in the current state as a traditional Pandas 2D dataframe converted to 2D numpy array, and the neural network should learn to select the row index of the best action given the features of each row.

So the input dataframe for one state could be something like:

df = pd.DataFrame({'user_rating':[10,7,3],'professional_rating':[10,4,1]})

Screenshot (18)

and it would learn to output index 0 since it has the highest rating.

But then the next state could have the highest rating in the 3rd row instead of the 1st and it would need to pick up on that. (Tried adding another screenshot of this but new users can only add 1).

It doesn’t seem to work when I flatten out the state and pass it in as a 1D array. Is there any other way of doing this in PyTorch?

Thanks so much!