Thank you Thomas.
I’m afraid my explanation was not enough and so confused you. Suppose grid_size=3 and only one object whose label is 3 is in grid [1,2] then the target is
[[0,0,0],
[0,0,0],
[0,3,0]]

The cross entropy loss is ce(output[:, 1, 2], [3]). I used nonzero to get all the indices of object existing grid(s) but that was a bit slow