GPU사용

CUDA함수

• 엔비디아 그래픽카드 지원

x = torch.cuda.FloatTensor(2,2)

결과
tensor([[0,0],[0,0]], device = 'cuda:0')

• cuda뒤의 숫자는 디바이스의 인덱스
• 0번 GPU 의미

다만, 서로 다른 장치에 올라가 있는 텐서끼리는 연산이 안됨.

1. nn.Module이 아닌곳에서의 cuda

• 복사가 됨

x = x.cuda()를 해도 gpu사용됨
x = x.cuda(device =1)로 인덱스 변경가능

1. nn.Module하에서 cuda()

• 텐서 이동

import torch.nn as nn
layer = nn.Linear(2,2)
layer.cuda(0)

CPU함수

x = x.cpu()