@@ -56,7 +56,7 @@ def build_dataset(args, trans_config, device, is_train=False):
dataset = COCODataset(
img_size=args.img_size,
data_dir=data_dir,
- image_set='train2017',
+ image_set='train2017' if is_train else 'val2017',
transform=train_transform,
trans_config=trans_config,
is_train=is_train