From cf505c65495f5cbbad0790c631677c99b308b740 Mon Sep 17 00:00:00 2001 From: whenyd Date: Fri, 13 Mar 2020 11:21:16 +0800 Subject: [PATCH] Set `drop_last=True` for val_loader Former-commit-id: adb1d7e9348a2707f38d9b86f57bb3ae2cbc2b73 --- train.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/train.py b/train.py index 5ac3f3b..90a60ae 100644 --- a/train.py +++ b/train.py @@ -35,7 +35,7 @@ def train_net(net, n_train = len(dataset) - n_val train, val = random_split(dataset, [n_train, n_val]) train_loader = DataLoader(train, batch_size=batch_size, shuffle=True, num_workers=8, pin_memory=True) - val_loader = DataLoader(val, batch_size=batch_size, shuffle=False, num_workers=8, pin_memory=True) + val_loader = DataLoader(val, batch_size=batch_size, shuffle=False, num_workers=8, pin_memory=True, drop_last=True) writer = SummaryWriter(comment=f'LR_{lr}_BS_{batch_size}_SCALE_{img_scale}') global_step = 0