4クラスの画像から各クラスごとに最も確信度が低い点を選ぶ

前提・実現したいこと

4クラスの画像から各クラスごとに最も確信度が低い点(のインデックス)を選ぶようなプログラムを作っています。
[bactch,channel,H,W]=[4,64,256,256]
4クラスの画像から最も確信度が低い点(のインデックス)を選ぶプログラムは作ることができ、このコードから各クラスごとに最も確信度が低い点(4クラスだから合計4点)を選ぶようなプログラムを作りたいです。

発生している問題

今のコードは最も確信度が低い点(最大確率が最も低い点)を選んできているだけなので、クラスごとにどうやって確信度が低い点を求めたらいいのか困っています。
y1,y2 = torch.min(y1, dim=0, keepdim=False, out=None)←ここで全てのピクセルにおける最大確率の最小値を出して、最後にargminでインデックスを取得しているのですが、これをクラスごとに取得するテクニックなどがあれば知りたいです。

該当のソースコード

python3ここに言語名を入力
1        #一次元配列にする (b,c,H,W) -> (b,c,H*W)
2        y = y.view(y.size(0), y.size(1), y.size(2) * y.size(3))  #shape[4, 4, 65536]
3        
4        #channel方向にsoftmaxをかけて確率にする
5        y = F.softmax(y, dim=1)  #shape[4, 4, 65536]
6        
7        #列方向の最大値を返す y1に値、y2にインデックスを取得する 65536がbatch個の行列
8        y1,y2 = torch.max(y, dim=0, keepdim=False, out=None)
9        print("y1", y1) 
10        print("y2", y2)
11        print("y1.shape", y1.shape) #[4, 65536]
12        #print(y2.shape) #[4, 65536]
13        
14        #列方向の最小値を返す y1に値、y2にインデックスを取得する 65536が1個の行列
15        y1,y2 = torch.min(y1, dim=0, keepdim=False, out=None)
16        
17        #最も確率が低い値のインデックスを取得する
18        y3 = torch.argmin(y1, dim=None, keepdim=False) # 最大で65536まで
19        print("y3_index", y3)  #55772