device.py 158 Bytes
Newer Older
1
2
3
4
5
6
# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.

import torch

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