How to write torch.device('cuda' if torch.cuda.is_available() else 'cpu') as a full if else statement?

Question

I'm a beginner to Pytorch and wanted to type this statement as a whole if else statement:-

torch.device('cuda' if torch.cuda.is_available() else 'cpu')

Can somebody help me?

Answer 1

Here is the code as a whole if-else statement:

torch.device('cuda' if torch.cuda.is_available() else 'cpu')

if torch.cuda.is_available():
    torch.device('cuda')
else:
    torch.device('cpu')

Since you probably want to store the device for later, you might want something like this instead:

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

if torch.cuda.is_available():
    device = torch.device('cuda')
else:
    device = torch.device('cpu')

---

Here is a post and discussion about the ternary operator in Python: https://stackoverflow.com/a/2802748/13985765

From that post:

value_when_true if condition else value_when_false