PytorchでGPUが使えているか確認する方法

まず、一番シンプルなGPUが使えているかの確認です。

import torch
print(torch.cuda.is_available())

TrueであればpytorchはあなたのGPUを使うことができます。
Falseなら使えていません。

使えていない場合あなたのPCにCUDAが正しくインストールされていない可能性があります.
CUDAはGPUを科学計算用に使用するためのライブラリキットで，インストールは必須です．
↓からインストールします．
CUDA Toolkit 10.2 Download | NVIDIA Developer

windowsでインストールする人も多いと思いますが，個人的にはUbuntu(Linux)を強くオススメします．

次に、任意の配列をcpu用、gpu用で変換して使い分けてみます。
gpu用の配列はgpu用の配列としか演算できないので、コードの中でうまくcpuとgpuの配列の変換をしていく必要があります。

import torch
array = torch.zeros(4) #長さ4の配列を定義（これはcpu配列です）
print(array)
#結果は tensor([0., 0., 0., 0.]) と表示されるはず

array_gpu = array.cuda()
print(array_gpu)
#結果は tensor([0., 0., 0., 0.], device='cuda:0') と表示されるはず

人の環境によってはcuda:0のところは単にcudaのみかもしれないです。
上のコードから、pytorchでは簡単にcpu配列、gpu配列を切り替えることができることがわかりましたね！
まとめると以下のようになります。

import torch
array = torch.zeros(4) #長さ4の配列を定義（これはcpu配列です）
array_gpu = array.cuda() # .cuda()をつけることで、cpu配列からgpu配列に変換
array_cpu = array_gpu.cpu() # .cpu()をつけることで、gpu配列からcpu配列に変換

次に、複数のGPUがある時に特定のGPUを使う方法です。
例えば、PCに2つのGPUが入っている時を例にします。
おすすめなのは、先にtorch.deviceを使う方法。

import torch
dev_gpu0 = torch.device('cuda:0')
dev_gpu1 = torch.device('cuda:1')

array = torch.zeros(4)
array_0 = array.to(dev_gpu0) #1個目のgpuを使う
array_1 = array.to(dev_gpu1) #2個目のgpuを使う

gpuのデバイス識別はcuda:0, cuda:1のようにすることで制御できます。
それをtorch.deviceでデバイスの変数を作り、配列の後に .to(device)を呼び出すことで、任意のgpuの配列に変換することができます。