张量跨设备迁移时触发CUDA RuntimeError问题求助

将Tensor迁移至CUDA设备，或从CUDA设备迁移至CPU时均触发错误，已检查Tensor的shape与dtype均无异常，请问可能是什么原因？

报错栈信息 #

/opt/conda/conda-bld/pytorch_1678402411778/work/aten/src/ATen/native/cuda/Indexing.cu:1146: 
indexSelectLargeIndex: block: [55,0,0], thread: [0,0,0] 断言 `srcIndex < srcSelectDimSize` 失败。
......
Traceback (最近的调用在最后):
  文件 "/home/xc/.conda/envs/molbart_new/lib/python3.10/site-packages/lightning/pytorch/trainer/call.py", 第44行, 在 _call_and_handle_interrupt
    return trainer_fn(*args, **kwargs)
  文件 "/home/xc/.conda/envs/molbart_new/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", 第580行, 在 _fit_impl
    self._run(model, ckpt_path=ckpt_path)
  文件 "/home/xc/.conda/envs/molbart_new/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", 第989行, 在 _run
    results = self._run_stage()
  文件 "/home/xc/.conda/envs/molbart_new/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", 第1033行, 在 _run_stage
    self._run_sanity_check()
  文件 "/home/xc/.conda/envs/molbart_new/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", 第1062行, 在 _run_sanity_check
    val_loop.run()
  文件 "/home/xc/.conda/envs/molbart_new/lib/python3.10/site-packages/lightning/pytorch/loops/utilities.py", 第182行, 在 _decorator
    return loop_run(self, *args, **kwargs)
  文件 "/home/xc/.conda/envs/molbart_new/lib/python3.10/site-packages/lightning/pytorch/loops/evaluation_loop.py", 第134行, 在 run
    self._evaluation_step(batch, batch_idx, dataloader_idx, dataloader_iter)
  文件 "/home/xc/.conda/envs/molbart_new/lib/python3.10/site-packages/lightning/pytorch/loops/evaluation_loop.py", 第391行, 在 _evaluation_step
    output = call._call_strategy_hook(trainer, hook_name, *step_args)
  文件 "/home/xc/.conda/envs/molbart_new/lib/python3.10/site-packages/lightning/pytorch/trainer/call.py", 第309行, 在 _call_strategy_hook
    output = fn(*args, **kwargs)
  文件 "/home/xc/.conda/envs/molbart_new/lib/python3.10/site-packages/lightning/pytorch/strategies/strategy.py", 第403行, 在 validation_step
    return self.lightning_module.validation_step(*args, **kwargs)
  文件 "/home/xc/xc_mol_seq/xc_work/seq_template/modules/model.py", 第1027行, 在 validation_step
    bs, logits, loss = self.forward(batch)
  文件 "/home/xc/xc_mol_seq/xc_work/seq_template/modules/model.py", 第934行, 在 forward
    indices, templates_candidates, templates_candidates_score = self.topk_candidates(
  文件 "/home/xc/xc_mol_seq/xc_work/seq_template/modules/model.py", 第1138行, 在 topk_candidates
    scores, indices = scores.cpu().detach().numpy(), indices.cpu().detach().numpy()
RuntimeError: CUDA错误：设备端断言触发
CUDA内核错误可能会在其他API调用处异步报告，因此下面的栈跟踪可能不准确。
调试时可考虑设置环境变量 CUDA_LAUNCH_BLOCKING=1。
编译时添加 `TORCH_USE_CUDA_DSA` 以启用设备端断言。

• 核心问题：报错中的srcIndex < srcSelectDimSize断言失败，说明代码存在索引越界操作，只是CUDA异步执行特性导致错误延迟到Tensor设备迁移时才抛出，并非迁移操作本身有问题。
• 排查与解决步骤：
  1. 开启同步调试定位根源：设置环境变量CUDA_LAUNCH_BLOCKING=1后运行代码，这样能获取准确的报错位置，找到真正触发索引越界的代码行，而非设备迁移的那一行。
  2. 检查indices张量的合法性：报错出现在indices.cpu()时，说明indices中的元素值超出了对应源张量的维度范围。比如用indices索引某张量时，索引值大于等于该张量对应维度的大小。
  3. 回溯topk_candidates的生成逻辑：检查生成indices的过程，比如topk操作的参数是否合理，是否存在计算错误导致索引值超出有效范围。
  4. 验证索引与张量维度的匹配性：确保所有索引操作中，索引值的范围和被索引张量的对应维度大小一致，例如被索引张量某维度大小为N，则索引值必须在[0, N-1]范围内。

内容的提问来源于stack exchange，提问作者Satoshi Finn