训练yolox-s出现RuntimeError: CUDA error: device-side assert triggered问题

尝试减少batch和input_size,没有效果，之前可以正常训练，现在报错：
/pytorch/aten/src/ATen/native/cuda/Loss.cu:102: operator(): block: [10,0,0], thread: [26,0,0] Assertion input_val >= zero && input_val <= one failed.
/pytorch/aten/src/ATen/native/cuda/Loss.cu:102: operator(): block: [10,0,0], thread: [27,0,0] Assertion input_val >= zero && input_val <= one failed.
/pytorch/aten/src/ATen/native/cuda/Loss.cu:102: operator(): block: [10,0,0], thread: [28,0,0] Assertion input_val >= zero && input_val <= one failed.
/pytorch/aten/src/ATen/native/cuda/Loss.cu:102: operator(): block: [10,0,0], thread: [29,0,0] Assertion input_val >= zero && input_val <= one failed.
/pytorch/aten/src/ATen/native/cuda/Loss.cu:102: operator(): block: [10,0,0], thread: [30,0,0] Assertion input_val >= zero && input_val <= one failed.
/pytorch/aten/src/ATen/native/cuda/Loss.cu:102: operator(): block: [10,0,0], thread: [31,0,0] Assertion input_val >= zero && input_val <= one failed.
2022-03-03 14:35:41 | ERROR | yolox.models.yolo_head:328 - OOM RuntimeError is raised due to the huge memory cost during label assignment. CPU mode is applied in this batch. If you want to avoid this issue, try to reduce the batch size or image size.
2022-03-03 14:35:41 | INFO | yolox.core.trainer:196 - Training of experiment is done and the best AP is 0.00
2022-03-03 14:35:41 | ERROR | yolox.core.launch:98 - An error has been caught in function 'launch', process 'MainProcess' (32), thread 'MainThread' (140133772625728):
Traceback (most recent call last):

File "/project/train/src_repo/YOLOX/tools/../yolox/models/yolo_head.py", line 322, in get_losses
imgs,
└

File "/usr/local/lib/python3.6/dist-packages/torch/autograd/grad_mode.py", line 26, in decorate_context
return func(*args, **kwargs)
│ │ └ {}
│ └
└ <function YOLOXHead.get_assignments at 0x7f731c7157b8>

File "/project/train/src_repo/YOLOX/tools/../yolox/models/yolo_head.py", line 505, in get_assignments
) = self.dynamic_k_matching(cost, pair_wise_ious, gt_classes, num_gt, fg_mask)
│ │ │ │ │ │ └
│ │ │ │ │ └ 2
│ │ │ │ └
│ │ │ └
│ │ └
│ └ <function YOLOXHead.dynamic_k_matching at 0x7f731c715950>
└ YOLOXHead(
(cls_convs): ModuleList(
(0): Sequential(
(0): BaseConv(
(conv): Conv2d(128, 128, kernel_size=...

File "/project/train/src_repo/YOLOX/tools/../yolox/models/yolo_head.py", line 616, in dynamic_k_matching
dynamic_ks = dynamic_ks.tolist()
│ └ <method 'tolist' of 'torch._C._TensorBase' objects>
└

RuntimeError: CUDA error: device-side assert triggered

During handling of the above exception, another exception occurred:

Traceback (most recent call last):

File "train.py", line 146, in
args=(exp, args),
│ └ Namespace(batch_size=8, cache=False, ckpt='/project/train/models/weight/yolox_s.pth', devices=0, dist_backend='nccl', dist_ur...
└ ╒═══════════════════╤════════════════════════════════════════════════════════════════════════════════════════════════════════...

File "/project/train/src_repo/YOLOX/tools/../yolox/core/launch.py", line 98, in launch
main_func(*args)
│ └ (╒═══════════════════╤═══════════════════════════════════════════════════════════════════════════════════════════════════════...
└ <function main at 0x7f731c747ae8>

File "train.py", line 124, in main
trainer.train()
│ └ <function Trainer.train at 0x7f72d48b2bf8>
└ <yolox.core.trainer.Trainer object at 0x7f731c7589e8>

File "/project/train/src_repo/YOLOX/tools/../yolox/core/trainer.py", line 74, in train
self.train_in_epoch()
│ └ <function Trainer.train_in_epoch at 0x7f72d48d3f28>
└ <yolox.core.trainer.Trainer object at 0x7f731c7589e8>

File "/project/train/src_repo/YOLOX/tools/../yolox/core/trainer.py", line 83, in train_in_epoch
self.train_in_iter()
│ └ <function Trainer.train_in_iter at 0x7f731c745950>
└ <yolox.core.trainer.Trainer object at 0x7f731c7589e8>

File "/project/train/src_repo/YOLOX/tools/../yolox/core/trainer.py", line 89, in train_in_iter
self.train_one_iter()
│ └ <function Trainer.train_one_iter at 0x7f731c7459d8>
└ <yolox.core.trainer.Trainer object at 0x7f731c7589e8>

File "/project/train/src_repo/YOLOX/tools/../yolox/core/trainer.py", line 103, in train_one_iter
outputs = self.model(inps, targets)
│ │ │ └
│ │ └
│ └ YOLOX(
│ (backbone): YOLOPAFPN(
│ (backbone): CSPDarknet(
│ (stem): Focus(
│ (conv): BaseConv(
│ (conv): ...
└ <yolox.core.trainer.Trainer object at 0x7f731c7589e8>

File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 727, in _call_impl
result = self.forward(*input, **kwargs)
│ │ │ └ {}
│ │ └
│ └ <function YOLOX.forward at 0x7f731c715c80>
└ YOLOX(
(backbone): YOLOPAFPN(
(backbone): CSPDarknet(
(stem): Focus(
(conv): BaseConv(
(conv): ...

File "/project/train/src_repo/YOLOX/tools/../yolox/models/yolox.py", line 35, in forward
fpn_outs, targets, x
│ │ └
│ └
└

File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 727, in _call_impl
result = self.forward(*input, **kwargs)
│ │ │ └ {}
│ │ └
│ └ <function YOLOXHead.forward at 0x7f731c715510>
└ YOLOXHead(
(cls_convs): ModuleList(
(0): Sequential(
(0): BaseConv(
(conv): Conv2d(128, 128, kernel_size=...

File "/project/train/src_repo/YOLOX/tools/../yolox/models/yolo_head.py", line 203, in forward
dtype=xin[0].dtype,
└

File "/project/train/src_repo/YOLOX/tools/../yolox/models/yolo_head.py", line 352, in get_losses
"cpu",

File "/usr/local/lib/python3.6/dist-packages/torch/autograd/grad_mode.py", line 26, in decorate_context
return func(*args, **kwargs)
│ │ └ {}
│ └
└ <function YOLOXHead.get_assignments at 0x7f731c7157b8>

File "/project/train/src_repo/YOLOX/tools/../yolox/models/yolo_head.py", line 446, in get_assignments
gt_bboxes_per_image = gt_bboxes_per_image.cpu().float()
│ └ <method 'cpu' of 'torch._C._TensorBase' objects>
└

RuntimeError: CUDA error: device-side assert triggered
terminate called after throwing an instance of 'c10::Error'
what(): CUDA error: device-side assert triggered

Start with yolox/models/yolo_head.py, especially get_losses, get_assignments, and dynamic_k_matching, then trace the training entry point in train.py. Reproduce the CUDA assertion and determine its underlying cause; done should include a verified fix and evidence that YOLOX training no longer fails for the reported case.