switch the GPU

[a1112]

I am unable to switch the GPU used for training by setting device: cuda:1, which is not conducive to tasks that involve training multiple models simultaneously.

    if torch.cuda.is_available():
        if args.distributed:
            device = 'cuda:%d' % args.local_rank
        else:
            device = 'cuda:0'
        torch.cuda.set_device(device)
    else:
        device = 'cpu'
    args.device = device

Is it possible to use a change similar to the one below?

    if torch.cuda.is_available():
        if args.distributed:
            device = 'cuda:%d' % args.local_rank
        else:
            device = args.device if args.device else 'cuda:0'
        torch.cuda.set_device(device)
    else:
        device = 'cpu'
    device = torch.device(device)

[rwightman]

@a1112 so yeah, that's possible but there isn't actually a device arg passed via train script right now, it assumes either cuda (if cuda is present) or cpu if not. Should probably fix so it can also support other device backends.

I get what you're hoping for and it's doable have to rethink the logic a bit, ends up a little more complicated than your proposal with the fallback options and supporting devices outside of cpu/cuda....