| 191 | |
| 192 | class DataModuleTemplateFromConfig(DataModuleFromConfig): |
| 193 | def __init__(self, batch_size, train=None, validation=None, test=None, |
| 194 | wrap=False, num_workers=None): |
| 195 | super().__init__(batch_size, train=train, validation=validation, |
| 196 | test=test, wrap=wrap, num_workers=num_workers) |
| 197 | if train is not None: |
| 198 | self.dataset_configs["train"] = train |
| 199 | self.train_dataloader = self._train_dataloader |
| 200 | |
| 201 | def _train_dataloader(self): |
| 202 | return DataLoader(self.datasets["train"], batch_size=self.batch_size*2, |