I’ve been working on a machine with multiple GPUs, and have needed to specific which GPU to run on. The docs 6 say that best practice is to use a torch.cuda.device context manager, so the following line appears at multiple points in my code:
with torch.cuda.device(f"cuda:{ARGS.gpu}"):
    <do-tensor-stuff>

Now I want to make this code able to run on a CPU. Understandably, this line is now giving the error
AssertionError: 
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

I can’t think of a neat solution. Obviously the following is not good because there would be so much duplicate code.
if torch.cuda.is_available():
    with torch.cuda.device(f"cuda:{ARGS.gpu}"):
        <do-tensor-stuff>
else:
    <do-tensor-stuff>