CUDA error: an Illegal memory access was encountered

[ali-fani-sd]

We are using DLRM model for personalization and we are getting CUDA error. By setting up CUDA_LAUNCH_BLOCKING flag and enabling cuda core dump, it pointed to two files where the issue might be happening
1: torchrec/distributed/embeddingbag.py: input_dist
2:torchrec/sparse/jagged_tensor.py: permute()

Some of our jaggedtensors are using weights, so when we debug the Jagged_tenosor.py we see mismatch in values(permuted length per key sum) and weights. Do you think that could be the root cause of CUDA error.

[TroyGarden]

these two functions (input_dist and permute) are used in daily production, so unlikely there's issue with the functions themselves. without enough context my wild guess is that the input KJT might contain incorrect data. please check our kjt validator.

@spmex