cuda, mps, cpu

[mattiaskjellsson]

It seems to me that there are assumptions of using CUDA? I'm very new to python, but would it make sense to do something like this
if torch.cuda.is_available(): torch.cuda.set_device(local_rank) elif torch.mps.is_available(): torch.mps.set_device(local_rank) else: torch.cpu.set_device(local_rank)

in places where currently

torch.cuda.set_device(local_rank)

is used, or have I misunderstood something?