fixed ground truth dataloader

MiscellaneousStuff · May 5, 2022 · f3732e0 · f3732e0
1 parent 51d3be6
commit f3732e0
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/train.py b/train.py
@@ -103,12 +103,12 @@ def get_dataloaders(
  train_loader = data.DataLoader(dataset=train_dataset,
  batch_size=hparams['batch_size'],
  shuffle=True,
- collate_fn=lambda x: data_processing(x, 'train'),
+ collate_fn=lambda x: data_processing(x, encoder, 'train'),
  **kwargs)
  test_loader = data.DataLoader(dataset=test_dataset,
  batch_size=hparams['batch_size'],
  shuffle=False,
- collate_fn=lambda x: data_processing(x, 'valid'),
+ collate_fn=lambda x: data_processing(x, encoder, 'valid'),
  **kwargs)
  return train_loader, test_loader