| 190 | return dataset |
| 191 | |
| 192 | def get_dataloader(self, type_path: str, batch_size: int, shuffle: bool = False) -> DataLoader: |
| 193 | dataset = self.get_dataset(type_path) |
| 194 | sampler = None |
| 195 | if self.hparams.sortish_sampler and type_path == "train": |
| 196 | assert self.hparams.gpus <= 1 # TODO: assert earlier |
| 197 | sampler = dataset.make_sortish_sampler(batch_size) |
| 198 | shuffle = False |
| 199 | |
| 200 | dataloader = DataLoader( |
| 201 | dataset, |
| 202 | batch_size=batch_size, |
| 203 | collate_fn=dataset.collate_fn, |
| 204 | shuffle=shuffle, |
| 205 | num_workers=self.num_workers, |
| 206 | sampler=sampler, |
| 207 | ) |
| 208 | return dataloader |
| 209 | |
| 210 | def train_dataloader(self) -> DataLoader: |
| 211 | dataloader = self.get_dataloader("train", batch_size=self.hparams.train_batch_size, shuffle=True) |