cuda 사용 가능 여부에 따라 gpu 또는 cpu로 설정하는 방법은 아래 코드를 사용하면 돼요.
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
