File size: 87 Bytes
8cd44cd
 
 
1
2
3
4
import torch

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