Pytorch torch 버전, GPU 사용 유무, 이름 확인, device 설정

import torch

# torch version, gpu 사용 여부, 이름 체크
print("torch version>>", torch.__version__)
print("GPU 사용 유무>>", torch.cuda.is_available())
print("GPU 이름>> ", torch.cuda.get_device_name(0))

'''
torch version>> 1.7.1+cu101
GPU 사용 유무>> True
GPU 이름>>  GeForce GTX 1650
'''

# device 설정. GPU 사용 가능 여부에 따라 device 정보 저장
device = "cuda:0" if torch.cuda.is_available() else "cpu"
print(device)
# cuda:0