파이토치 GPU 인식 확인하기

torch-gpu-test.py 파일을 생성하여 아래 내용을 복붙한다.

import torch

def main():
    print(torch.cuda.is_available())
    print(torch.cuda.get_device_name(0))
    print(torch.cuda.device_count())

if __name__ == '__main__':
    main()

 

파이썬으로 위 프로그램을 실행하였을 때, 디바이스 이름과 개수가 의도한 바와 일치하면 된다.

 python torch-gpu-test.py

True
Tesla V100S-PCIE-32GB
2