(pytorch) 텐서를 CPU와 GPU 간에 이동시키기

릿카。 2024. 9. 4. 13:13

pytorch의 tensor는 넘파이의 ndarray와 유사하게 구현되었다고 하며, 차이점은 (1) GPU에서 연산을 수행할 수 있으며 (2) Autograd 계산에 최적화되어있다는 것이다.

if torch.cuda_is_available():
	tensor = tensor.to('cuda')

이런식으로 GPU로 옮겨주고,

tensor = tensor.to('cpu')

이렇게 CPU로 옮겨준다.