파이토치 cuda tensor

Yoonsnee·2022년 12월 23일
0

cuda tensor

변수명.device

: 현재 텐서가 어디에(cuda인지 cpu인지) 위치해있는지 알 수 있음

a = torch.tensor([[1,2,3,4],
                 [5,6,7,8]])
a.device 

# 결과값
# device(type='cpu') # 텐서가 cpu에 위치해있음

.to('cuda')

: 텐서의 위치를 cpu -> cuda

.to('cpu')

: 텐서의 위치를 cuda -> cpu

print(a.to('cuda'))

# 결과값
# tensor([[1, 2, 3, 4],
#       [5, 6, 7, 8]], device='cuda:0') # cpu에 위치해있던 텐서가 'cuda'로 옮겨진 것을 알 수 있음.

torch.cuda.is_available()

: cuda를 사용여부를 확인 할 수 있음.

b = torch.tensor([1, 2, 3, 4])

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

b = b.to(device)
print(b.device)

# 결과값
# cuda:0
profile
윤쓰네뽀끼

0개의 댓글