CUDA Initialization Warning on Docker when not passing in gpu

[NanoCode012]

🐛 Bug

train.py gives off CUDA warning when trying to use docker with cpu only.

To Reproduce (REQUIRED)

1. Start latest docker container without gpus
2. Run python train.py --device cpu --epoch 2

Output:

root@641d4b75809c:/usr/src/app# python train.py --device cpu --epoch 2
/opt/conda/lib/python3.8/site-packages/torch/cuda/__init__.py:52: UserWarning: CUDA initialization: Found no NVIDIA driver on your system. Please check that you have an NVIDIA GPU and installed a driver from http://www.nvidia.com/Download/index.aspx (Triggered internally at  ../c10/cuda/CUDAFunctions.cpp:100.)
  return torch._C._cuda_getDeviceCount() > 0
YOLOv5 v4.0-6-g6ab5895 torch 1.8.0a0+1606899 CPU

Program runs as expected afterwards.

Expected behavior

No warning given

Environment

• OS: Docker
• GPU: None

Additional context

Offending line is

yolov5/utils/torch_utils.py

Line 64 in 6ab5895

cuda = torch.cuda.is_available() and not cpu

A simple fix would be to change to

cuda = not cpu and torch.cuda.is_available()

PR that changed this: #1826

I checked the CI and noticed that they don't receive this warning which is confusing since they should not have any GPUs too.

[glenn-jocher]

@NanoCode012 that's odd! Your recommended change doesn't hurt anything though, so I've started a PR to implement it.

[NanoCode012]

Thanks!