Thanks for your work!!
I was trying to re-execute your code and found out the following issue.

And it returns an error message.

I think it has to do with the working environment(I'm using Windows 10).
I post the link for your information (https://stackoverflow.com/questions/56360644/pytorch-runtimeerror-expected-tensor-for-argument-1-indices-to-have-scalar-t)
Any feedback or modifications will be very appreciated!!