mirror of
https://github.com/Laurent2916/REVA-QCAV.git
synced 2024-11-09 15:02:03 +00:00
Set drop_last=True
for val_loader
Former-commit-id: adb1d7e9348a2707f38d9b86f57bb3ae2cbc2b73
This commit is contained in:
parent
d292e8c6cd
commit
cf505c6549
2
train.py
2
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
|
||||
|
|
Loading…
Reference in a new issue