CNNの画像を二次元配列から三次元配列へ変換

Question

[Python CNN Grad-CAM]
CNNの画像にノイズを加えたものの特徴マップをGrad-CAMをを用いて可視化することを目標としています。
ノイズを加える前の画像は可視化することができましたが、ノイズを加えた後の画像は、同じプログラムでは実行できませんでした。

ノイズなしの場合は、[1,28,28]の三次元、ありの場合は[28,28]の二次元配列です。
以下がエラー箇所のプログラムです。
画像データの引数はdataです。

data.shape[2]

でエラーが発生します。

二次元配列のデータはそのままで、軸を増やせば解決できると思うのですが、どうすれば次元を増やせるでしょうか。
ご教授よろしくおねがいします。

以下の二行目がエラー部です。

data, label = trainset[0]
input_data = data.view(1, data.shape[0], data.shape[1], data.shape[2]).to(device)

# モデルで予測
# confidence: 信頼度、 predicted: 予測ラベル
# (今回は、torch.max(, 1)でTop1の情報のみを取得
output = nn.Softmax(dim=1)(model(input_data))
confidence, predicted = torch.max(output.data, 1)

# GradCAMオブジェクト生成
gcam = GradCAM(model=model)

# 画像を順伝搬
_ = gcam.forward(input_data)

# 予測ラベルを元に逆伝搬し、勾配を計算
single_predicted = predicted.view(1, predicted.shape[0]).to(device)
gcam.backward(ids=single_predicted)

# 指定の層の勾配マップを取得
#(各層の名前は、print(model)で参照可能)
regions = gcam.generate("conv1")

# tensorを、numpy に変換
raw_image = input_data[0].to('cpu').detach().numpy().copy()
raw_image = raw_image.transpose((1, 2, 0))
raw_image = ((raw_image * 0.5) + 0.5) * 255.0
raw_image = raw_image.astype(np.uint8)

# 各画像 出力
output = out_gradcam(
    gcam=regions[0, 0],
    raw_image=raw_image
)

## GradCAM
plt.imshow(output, cmap="jet")
plt.colorbar()

Answer 1

>>> data = data.reshape(1,28,28)
>>> data.shape
(1, 28, 28)

です。