Based on the error message, it seems that you are using the GPU:
/pytorch/aten/src/THC/THCTensorIndex.cu:361: void indexSelectLargeIndex(TensorInfo<T, IndexType>, TensorInfo<T, IndexT
ype>, TensorInfo<long, IndexType>, int, int, IndexType, IndexType, long) [with T = float, IndexType = unsigned int, Ds
tDim = 2, SrcDim = 2, IdxDim = -2, IndexIsMajor = true]: block: [145,0,0], thread: [95,0,0] Assertion srcIndex < srcS electDimSize failed.

Which PyTorch version are you using? Could you update to the nightly binaries and rerun the code, please?