Pytorch 他クラス分類 各ラベルの確率で出力

PytorchでCifar-100を使用して20種類の画像分類を行っています。
VGG-16と学習させたモデルのパラメータを読み込んで推論を行う部分にて質問があります。

以下は確率の高いラベル一つだけを選んでますが、アウトプットとしてはsample_submit.csvのように画像ファイルに対して20個各ラベルの確率を出力したいと考えています。以下のargmaxのところで一つのラベルを選択しているのは把握しておりますが、各ラベルの確率として出力するにはどのように改変するには至らず、ご教示いただけませんでしょうか？
sample_submit.csv

# 学習済みのVGG-16モデルをロード # VGG-16モデルのインスタンスを生成 use_pretrained = True net = models.vgg16(pretrained=use_pretrained) # VGG16の最後の出力層の出力ユニットを付け替える net.classifier[6] = nn.Linear(in_features=4096, out_features=20) # PyTorchのネットワークパラメータのロード load_path = '/content/drive/MyDrive/data/20labels/output/weights_fine_tunig.pth' load_weights = torch.load(load_path) net.load_state_dict(load_weights) net.eval() # 推論 def predict(net, test_dataloader): device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') print('使用デバイス: ', device) # ネットワークをGPUへ net.to(device) # ネットワークの順伝搬および誤差関数の計算手法がある程度一定であれば、GPUの計算を高速化 torch.backends.cudnn.benchmark = True net.eval() preds = [] # データ ローダーからミニバッチを取り出すループ for inputs, _ in tqdm(test_dataloader): # GPUが使えるならGPUにデータを送る inputs = inputs.to(device) # 順伝搬(forward)計算 with torch.no_grad(): outputs = net(inputs) # NNを通して出力計算 preds += [int(output.argmax()) for output in outputs] return preds preds = predict(net, test_dataloader) test_df['label_id'] = preds save_path = '/content/drive/MyDrive/data/20labels/output/submission.csv' test_df.to_csv(save_path, index=False,header=False)