CUDA问题

[acheng777]

报错信息如下：
root@autodl-container-7154118f52-44680f8e:~/edgeyolo-main# python train.py --cfg ./params/train/train_coco.yaml
Traceback (most recent call last):
File "train.py", line 16, in
train("DEFAULT" if args.default else args.cfg)
File "/root/edgeyolo-main/edgeyolo/train/launch_train.py", line 101, in launch
mp.start_processes(
File "/root/miniconda3/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 188, in start_processes
while not context.join():
File "/root/miniconda3/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 150, in join
raise ProcessRaisedException(msg, error_index, failed_process.pid)
torch.multiprocessing.spawn.ProcessRaisedException:

-- Process 1 terminated with the following error:
Traceback (most recent call last):
File "/root/miniconda3/lib/python3.8/site-packages/torch/multiprocessing/spawn.py", line 59, in _wrap
fn(i, *args)
File "/root/edgeyolo-main/edgeyolo/train/launch_train.py", line 50, in train_single
torch.cuda.set_device(device)
File "/root/miniconda3/lib/python3.8/site-packages/torch/cuda/init.py", line 261, in set_device
torch._C._cuda_setDevice(device)
RuntimeError: CUDA error: invalid device ordinal
请问怎么解决？

[acheng777]

为什么没人回复我，呜呜呜

[WainBK]

你可以改一下train_coco.yaml中的 device