15 lines
299 B
Python
15 lines
299 B
Python
##
|
|
import torch
|
|
from torchvision import datasets, transforms, models
|
|
|
|
def get_device():
|
|
|
|
if torch.cuda.is_available():
|
|
# clear cach to boost memory
|
|
# for new round
|
|
torch.cuda.empty_cache()
|
|
return torch.device("cuda")
|
|
else:
|
|
return torch.device("cpu")
|
|
|