The device assert is the “hard error” option for CUDA kernels: It invalidates the CUDA context, which for practical purposes means you need to restart PyTorch. The problem here is that there isn’t any good method to report errors due to the asynchronous nature of CUDA. To find out which kernel is the one erroring, either insert torch.cuda.synchronize() between lines or run with the environment variable CUDA_LANCH_BLOCKING=1 set.
My guess would be torch.multinomial. For example, it is allergic to all-zero probability tensors (and because it returns integers cannot use NaN).
Best regards
Thomas