def _get_data_loader()

in torchrecipes/text/doc_classification/datamodule/doc_classification.py [0:0]

• 19 lines of code
• 1 McCabe index (conditional complexity)

    def _get_data_loader(self, dataset: IterDataPipe[Tuple[str, str]]) -> DataLoader:
        dataset = dataset.batch(self.batch_size).rows2columnar(["text", "label"])
        dataset = dataset.map(self.transform)
        dataset = dataset.map(
            lambda x: {
                **x,
                "label_ids": to_tensor(self.label_transform(x["label"])),
            }
        )
        dataset = dataset.add_index()

        return DataLoader(
            dataset,
            batch_size=None,
            shuffle=False,
            num_workers=self.num_workers,
            drop_last=self.drop_last,
            pin_memory=self.pin_memory,
            worker_init_fn=worker_init_fn,
        )