# Buggy code
def select_device(device="", batch_size=None):
    # device = 'cpu' or '0' or '0,1,2,3'
    cpu_request = device.lower() == "cpu"
    if device and not cpu_request:  # if device requested other than 'cpu'
        os.environ["CUDA_VISIBLE_DEVICES"] = device  # set environment variable
        assert torch.cuda.is_available(), (
            "CUDA unavailable, invalid device %s requested" % device
        )  # check availablity

    cuda = False if cpu_request else torch.cuda.is_available()
    if cuda:
        c = 1024**2  # bytes to MB
        ng = torch.cuda.device_count()
        if (
            ng > 1 and batch_size
        ):  # check that batch_size is compatible with device_count
            assert batch_size % ng == 0, (
                "batch-size %g not multiple of GPU count %g" % (batch_size, ng)
            )
        x = [torch.cuda.get_device_properties(i) for i in range(ng)]
        s = f"Using torch {torch.__version__} "
        for i in range(0, ng):
            if i == 1:
                s = " " * len(s)
            logger.info(
                "%sCUDA:%g (%s, %dMB)" % (s, i, x[i].name, x[i].total_memory / c)
            )
    else:
        logger.info(f"Using torch {torch.__version__} CPU")

    logger.info("")  # skip a line
    return torch.device("cuda:0" if cuda else "cpu")


# Traceback (RuntimeError)
$ python train.py --img 640 --batch 16 --epochs 3 --data coco128.yaml --weights yolov5x.pt --device "2"
Using torch 1.7.1 CUDA:0 (GeForce RTX 2080 Ti, 11019MB)

Namespace(adam=False, batch_size=16, bucket='', cache_images=False, cfg='', data='./data/coco128.yaml', device='2', epochs=3, evolve=False, exist_ok=False, global_rank=-1, hyp='data/hyp.scratch.yaml', image_weights=False, img_size=[640, 640], local_rank=-1, log_artifacts=False, log_imgs=16, multi_scale=False, name='exp', noautoanchor=False, nosave=False, notest=False, project='runs/train', rect=False, resume=False, save_dir='runs/train/exp22', single_cls=False, sync_bn=False, total_batch_size=16, weights='yolov5x.pt', workers=8, world_size=1)
Start Tensorboard with "tensorboard --logdir runs/train", view at http://localhost:6006/
Hyperparameters {'lr0': 0.01, 'lrf': 0.2, 'momentum': 0.937, 'weight_decay': 0.0005, 'warmup_epochs': 3.0, 'warmup_momentum': 0.8, 'warmup_bias_lr': 0.1, 'box': 0.05, 'cls': 0.5, 'cls_pw': 1.0, 'obj': 1.0, 'obj_pw': 1.0, 'iou_t': 0.2, 'anchor_t': 4.0, 'fl_gamma': 0.0, 'hsv_h': 0.015, 'hsv_s': 0.7, 'hsv_v': 0.4, 'degrees': 0.0, 'translate': 0.1, 'scale': 0.5, 'shear': 0.0, 'perspective': 0.0, 'flipud': 0.0, 'fliplr': 0.5, 'mosaic': 1.0, 'mixup': 0.0}

……
Model Summary: 607 layers, 88965245 parameters, 88965245 gradients, 221.7 GFLOPS

Transferred 802/802 items from yolov5x.pt
Optimizer groups: 134 .bias, 142 conv.weight, 131 other
Scanning '../coco128/labels/train2017.cache' for images and labels... 128 found, 0 missing, 2 empty, 0 corrupted: 100%|███████████████████████████████████████████████████████████████████| 128/128 [00:00<?, ?it/s]
Scanning '../coco128/labels/train2017.cache' for images and labels... 128 found, 0 missing, 2 empty, 0 corrupted: 100%|███████████████████████████████████████████████████████████████████| 128/128 [00:00<?, ?it/s]
Scanning '../coco128/labels/train2017.cache' for images and labels... 128 found, 0 missing, 2 empty, 0 corrupted: 100%|███████████████████████████████████████████████████████████████████| 128/128 [00:00<?, ?it/s]
Scanning '../coco128/labels/train2017.cache' for images and labels... 128 found, 0 missing, 2 empty, 0 corrupted: 100%|███████████████████████████████████████████████████████████████████| 128/128 [00:00<?, ?it/s]
Scanning '../coco128/labels/train2017.cache' for images and labels... 128 found, 0 missing, 2 empty, 0 corrupted: 100%|███████████████████████████████████████████████████████████████████| 128/128 [00:00<?, ?it/s]
Scanning '../coco128/labels/train2017.cache' for images and labels... 128 found, 0 missing, 2 empty, 0 corrupted: 100%|███████████████████████████████████████████████████████████████████| 128/128 [00:00<?, ?it/s]
Scanning '../coco128/labels/train2017.cache' for images and labels... 128 found, 0 missing, 2 empty, 0 corrupted: 100%|███████████████████████████████████████████████████████████████████| 128/128 [00:00<?, ?it/s]
Plotting labels...
Scanning '../coco128/labels/train2017.cache' for images and labels... 128 found, 0 missing, 2 empty, 0 corrupted: 100%|███████████████████████████████████████████████████████████████████| 128/128 [00:01<?, ?it/s]
Scanning '../coco128/labels/train2017.cache' for images and labels... 128 found, 0 missing, 2 empty, 0 corrupted: 100%|███████████████████████████████████████████████████████████████████| 128/128 [00:01<?, ?it/s]
Scanning '../coco128/labels/train2017.cache' for images and labels... 128 found, 0 missing, 2 empty, 0 corrupted: 100%|███████████████████████████████████████████████████████████████████| 128/128 [00:02<?, ?it/s]
Scanning '../coco128/labels/train2017.cache' for images and labels... 128 found, 0 missing, 2 empty, 0 corrupted: 100%|███████████████████████████████████████████████████████████████████| 128/128 [00:02<?, ?it/s]
Scanning '../coco128/labels/train2017.cache' for images and labels... 128 found, 0 missing, 2 empty, 0 corrupted: 100%|███████████████████████████████████████████████████████████████████| 128/128 [00:02<?, ?it/s]
Scanning '../coco128/labels/train2017.cache' for images and labels... 128 found, 0 missing, 2 empty, 0 corrupted: 100%|███████████████████████████████████████████████████████████████████| 128/128 [00:02<?, ?it/s]
Scanning '../coco128/labels/train2017.cache' for images and labels... 128 found, 0 missing, 2 empty, 0 corrupted: 100%|███████████████████████████████████████████████████████████████████| 128/128 [00:02<?, ?it/s]

Analyzing anchors... anchors/target = 4.26, Best Possible Recall (BPR) = 0.9946
Image sizes 640 train, 640 test
Using 8 dataloader workers
Logging results to runs/train/exp22
Starting training for 3 epochs...

Epoch   gpu_mem       box       obj       cls     total   targets  img_size
0%|                                                                                                                                                                                         | 0/8 [00:01<?, ?it/s]
Traceback (most recent call last):
File "train.py", line 512, in <module>
train(hyp, opt, device, tb_writer, wandb)
File "train.py", line 289, in train
pred = model(imgs)  # forward
File "/home/elfin/anaconda3/envs/yolov5/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
result = self.forward(*input, **kwargs)
File "/home/elfin/utils/yolov5-master/models/yolo.py", line 122, in forward
return self.forward_once(x, profile)  # single-scale inference, train
File "/home/elfin/utils/yolov5-master/models/yolo.py", line 138, in forward_once
x = m(x)  # run
File "/home/elfin/anaconda3/envs/yolov5/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
result = self.forward(*input, **kwargs)
File "/home/elfin/utils/yolov5-master/models/common.py", line 119, in forward
return torch.cat(x, self.d)
RuntimeError: CUDA out of memory. Tried to allocate 126.00 MiB (GPU 0; 10.76 GiB total capacity; 9.69 GiB already allocated; 11.56 MiB free; 9.74 GiB reserved in total by PyTorch)

# Fixed code
def select_device(device="", batch_size=None):
    # device = 'cpu' or '0' or '0,1,2,3'
    cpu_request = device.lower() == "cpu"
    if device and not cpu_request:  # if device requested other than 'cpu'
        os.environ["CUDA_VISIBLE_DEVICES"] = device  # set environment variable
        assert torch.cuda.is_available(), (
            f"CUDA unavailable, invalid device {device} requested"
        )  # check availablity

    cuda = False if cpu_request else torch.cuda.is_available()
    if cuda:
        c = 1024**2  # bytes to MB
        ng = torch.cuda.device_count()
        if (
            ng > 1 and batch_size
        ):  # check that batch_size is compatible with device_count
            assert batch_size % ng == 0, (
                f"batch-size {batch_size} not multiple of GPU count {ng}"
            )
        x = [torch.cuda.get_device_properties(i) for i in range(ng)]
        s = f"Using torch {torch.__version__} "
        for i, d in enumerate((device or "0").split(",")):
            if i == 1:
                s = " " * len(s)
            logger.info(f"{s}CUDA:{d} ({x[i].name}, {x[i].total_memory / c}MB)")
    else:
        logger.info(f"Using torch {torch.__version__} CPU")

    logger.info("")  # skip a line
    return torch.device("cuda:0" if cuda else "cpu")