【Pytorch】CUDA error: device-side assert triggered で発生したエラー

エラー

 CUDA error: device-side assert triggered 

ビビりました。
単刀直入によるとネットワークに入力したデータの形式が間違っている場合に、
発生するエラーだそうです。

やろうとしたこと

CIFAR10の画像から2クラスのみを取得して、二値分類を試そうとした。

調査

Jupyter Notebookで一行ずつ実装していたのですが、
一度、CUDA error: device-side assert triggeredが発生した場合、
KernelをRestartしないといけないようです。

問題のコードは以下の順伝播の計算の段階です。

for i , (images, labels) in enumerate(train_loader):
        
        # imageを一次元化→gpu
        # images, labels = images.view(-1, 32*32*3).to(device), labels.to(device)
        images, labels = images.to(device), labels.to(device)
        # Optimizerを0にリセット
        optimizer.zero_grad()
        # 順伝搬の計算
        features, outputs = net(images)

原因

CIFAR10の画像から取ってきた2クラスのラベル番号の振り分け。
自分が取得したクラスが5番と7番のクラスであったため、
このクラス番号を0と1に振り分けてやらなきゃいけません。
この処理を忘れていたので、上記のエラーが発生しました。

def label_classification(phase='train'):
    dog_indices, deer_indices = [], []
    for i in range(len(datasets[phase])):
        current_dataset = datasets[phase][i]
        # 犬画像の場合
        if current_dataset[1] == dog_idx:
            # 追加分(犬画像の場合、ラベル番号は0に)
            l = list(current_dataset)
            l[1] = 0
            current_dataset = tuple(l)
            dog_indices.append(current_dataset)
        # 鹿画像の場合
        elif current_dataset[1] == deer_idx:
            # 追加分(鹿画像の場合、ラベル番号は1に)
            l = list(current_dataset)
            l[1] = 1
            current_dataset = tuple(l)
            deer_indices.append(current_dataset)
    return dog_indices, deer_indices

これで正常に学習が完了しました。

結論

CUDA error: device-side assert triggered は個人によって対処法が違うので、
ネットでググるより仮定を立てて、解決した方が良い。