TypeError: __init__() got an unexpected keyword argument 'pin_memory'

[sooftware]

Hi! Thank you for your grrreat project. I encounter error log liker below:

• Error log

File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/data_loading.py", line 409, in reset_val_dataloader
    self.num_val_batches, self.val_dataloaders = self._reset_eval_dataloader(model, 'val')
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/data_loading.py", line 358, in _reset_eval_dataloader
    self.auto_add_sampler(dl, shuffle=False, mode=self.state.stage) for dl in dataloaders if dl is not None
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/data_loading.py", line 358, in <listcomp>
    self.auto_add_sampler(dl, shuffle=False, mode=self.state.stage) for dl in dataloaders if dl is not None
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/data_loading.py", line 142, in auto_add_sampler
    dataloader = self.replace_sampler(dataloader, sampler, mode=mode)
  File "/opt/conda/lib/python3.7/site-packages/pytorch_lightning/trainer/data_loading.py", line 213, in replace_sampler
    dataloader = type(dataloader)(**dl_args)
TypeError: __init__() got an unexpected keyword argument 'pin_memory'

• My code sample:

class AudioDataLoader(DataLoader):
    """ Audio Data Loader """
    def __init__(
            self,
            dataset: torch.utils.data.Dataset,
            num_workers: int,
            batch_sampler: torch.utils.data.sampler.Sampler,
    ) -> None:
        super(AudioDataLoader, self).__init__(dataset=dataset, num_workers=num_workers, batch_sampler=batch_sampler)
        self.collate_fn = _collate_fn


class BucketingSampler(Sampler):
    """ Samples batches assuming they are in order of size to batch similarly sized samples together. """
    def __init__(self, data_source, batch_size: int = 32, drop_last: bool = False) -> None:
        super(BucketingSampler, self).__init__(data_source)
        self.batch_size = batch_size
        self.data_source = data_source
        ids = list(range(0, len(data_source)))
        self.bins = [ids[i:i + batch_size] for i in range(0, len(ids), batch_size)]
        self.drop_last = drop_last

    def __iter__(self):
        for ids in self.bins:
            np.random.shuffle(ids)
            yield ids

    def __len__(self):
        return len(self.bins)

    def shuffle(self, epoch):
        np.random.shuffle(self.bins)

class LightningCustomDataModule(pl.LightningDataModule):
    def train_dataloader(self) -> DataLoader:
        train_sampler = BucketingSampler(self.dataset['train'], batch_size=self.batch_size)
        return AudioDataLoader(
            dataset=self.dataset['train'],
            num_workers=self.num_workers,
            batch_sampler=train_sampler,
        )

Please let me know if you have any doubts.