fig, axis = plt.subplots() の表示について

Question

実現したいこと

pythonでMNISTを使ってCNNの勉強をしています。

予測値が0となっているときのみ画像をプロットさせて、5×5で25個の異なる0を表示させたいのですが、それ以外の時もプロットされてしまい、結果下の画像のようになってしまいます。

該当のソースコード

Y_pred = model.predict(X_val)
Y_pred_classes = np.argmax(Y_pred, axis = 1) 
Y_true = np.argmax(Y_val, axis = 1) 
X_test2 = X_val.reshape(X_val.shape[0], 28, 28)

fig, axis = plt.subplots(5, 5, figsize=(4, 4))
for i, ax in enumerate(axis.flat):
    if Y_pred[i].argmax() == 0:
        ax.set_xticklabels([])
        ax.set_yticklabels([])
        ax.axis('off')
        ax.imshow(X_test2[i], cmap='binary')
        ax.set(title = f"True {Y_true[i]} \n Pred {Y_pred[i].argmax()}")
plt.show()

アドバイス等、頂けますと幸いです。
また、質問の仕方等、至らない箇所があれば指摘頂けるとうれしいです。
よろしくお願い致します。