Pytorch Datasetの作り方 train(valid)とtestで元にするファイル形式が異なる場合

PytorchでCifar-100を使用して20種類の画像分類を行っています。
自作Datasetの作り方の部分で質問があります。
Datasetを作成にあたりtrainとvalidとtestに分けてデータセットを作成します。trainとvalidは画像ファイル名と対応するラベルが記されたtrain_master.tsvを使用してdic形式にまとめた後に下に記載のLabelsDatasetでtrainとvalidのデータセットを作ります。
一方で、testに関しては同様な形で下の画像のようなsample_submit.csv より作成しようとしましたが、こちらは画像ファイル名と20個のラベルの確率を記したファイルのため同じようにできません。どのように修正すべきかご教示いただけませんでしょうか？
sample_submit.csv

# ライブラリ読み込み
import os
import glob
import random
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
import json
from PIL import Image
from tqdm import tqdm
import matplotlib.pyplot as plt
get_ipython().run_line_magic('matplotlib', 'inline')

import torch
import torch.nn as nn
import torch.optim as optim
import torch.utils.data as data
import torchvision
from torchvision import models, transforms

torch.manual_seed(1234)
np.random.seed(1234)
random.seed(1234)

# 前処理クラス定義
class ImageTransform():
    """
    画像の前処理クラス。訓練時、検証時で異なる動作をする
    """
    
    def __init__(self, resize, mean, std):
        self.data_transform = {
            'train': transforms.Compose([
                transforms.RandomResizedCrop(resize, scale=(0.5, 1.0)),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                transforms.Normalize(mean, std)
            ]),
            'val': transforms.Compose([
                transforms.Resize(resize),
                transforms.CenterCrop(resize),
                transforms.ToTensor(),
                transforms.Normalize(mean, std)
            ]),
            'test': transforms.Compose([
                transforms.Resize(resize),
                transforms.CenterCrop(resize),
                transforms.ToTensor(),
                transforms.Normalize(mean, std)
            ])
        }

         def __call__(self, img, phase='train'):
        return self.data_transform[phase](img)

        """
        Parameters
        ----------
        phase : 'train' or 'val'
            前処理のモードを指定。
        """

size = 224
mean = (0.485, 0.456, 0.406)
std = (0.229, 0.224, 0.225)

# 学習用、検証用にデータを分割
label_df = pd.read_csv('/content/drive/MyDrive/data/20labels/input/train_master.tsv',sep="\t")

test_size = 0.2
train_df, val_df = train_test_split(label_df, test_size=test_size, random_state=42, stratify=label_df['label_id'])
print(train_df.shape, val_df.shape)

#trainとvalに分けdic形式でまとめる
train_label_dic = dict(zip(train_df['file_name'], train_df['label_id']))
val_label_dic = dict(zip(val_df['file_name'], val_df['label_id']))

# Dataset作成
class LabelsDataset(data.Dataset):
    def __init__(self, label_dic, root_dir, transform=None, phase='train'):
        self.label_dic = label_dic
        self.root_dir = root_dir
        self.images = list(label_dic.keys())
        self.transform = transform
        self.phase = phase

    #__len__() : クラスインスタンスにlen()を使った時に呼ばれる関数    
    def __len__(self):
        return len(self.images)
    
    #__getitem__() : クラスインスタンスの要素を参照するときに呼ばれる関数
    def __getitem__(self, idx):
        filename = self.images[idx]
        img = Image.open(os.path.join(self.root_dir, filename))
        img_transformed = self.transform(img, self.phase)
        label = self.label_dic[filename]
        return img_transformed, label

train_dataset = LabelsDataset(train_label_dic, '/content/data/train/', ImageTransform(size, mean, std), 'train')
val_dataset = LabelsDataset(val_label_dic, '/content/data/train/', ImageTransform(size, mean, std), 'val')