Hi,
As you can see in the warning, the indices that you use for indexing are not in the proper range: they are larger than the size of that dimension. You should make sure your indices are correct.
I would guess this is most likely in your loss and the wrong indices are the labels.
Numpy does not support GPU. So you will have to convert your code to pytorch if you want to be able to use GPU.