Yet-another-efficientdet-pytorch: CUDA error: invalid device ordinal

[singlautsav]

I am running a custom training with product checkout dataset (coco)
I am using google colab with GPU and the following is the error

CUDA error: invalid device ordinal
0% 11/6717 [00:24<1:31:06, 1.23it/s][Error] Traceback (most recent call last):
File "train.py", line 227, in train
cls_loss, reg_loss = model(imgs, annot, obj_list=params.obj_list)
File "/usr/local/lib/python3.6/dist-packages/torch/nn/modules/module.py", line 532, in __call__
result = self.forward(*input, **kwargs)
File "/usr/local/lib/python3.6/dist-packages/torch/nn/parallel/data_parallel.py", line 148, in forward
inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)
File "/content/Yet-Another-EfficientDet-Pytorch/utils/utils.py", line 199, in scatter
for device_idx in range(len(devices))], \
File "/content/Yet-Another-EfficientDet-Pytorch/utils/utils.py", line 199, in
for device_idx in range(len(devices))], \
RuntimeError: CUDA error: invalid device ordinal

Training command that I have used for the same is

!python train.py -c 1 -p productDataset --batch_size 8 --lr 1e-5

[zylo117]

change num_gpu=1 in project.yml

[singlautsav]

I changed line 194 in utils.py to devices = ['cuda:0']
and yes your suggestion would have worked the same way

Thank you for the response I was going to close the issue :)