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)
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?