Sorry, ignore what I said- I misunderstood your code.
You are iterating through the DataLoader, but there is a bug in our _WrappedDataLoader where we return a generator instead of a tuple as intended: ray/torch.py at master · ray-project/ray · GitHub. So indexing into the batch will fail. I will make a PR to fix this, thanks for pointing this out!