Is your code working fine on the CPU?
If the error is raised in the first epoch, if might not take too long to run it on the CPU.
Based on the stack trace I’m not sure, where the failing indexing operation is used.