def create_data_loader(self, dataset, batch_size=32, mode="train", **kwargs):
        """
        Method to create a data loader for a given Tensor dataset.

        Args:
            mode(str): Mode for creating data loader. Could be train or test.
            dataset(torch.utils.data.Dataset): A Tensor dataset.
            batch_size(int): Batch size.

        Returns:
            torch.utils.data.DataLoader: A torch data loader to the given dataset.

        """

        if mode == "test":
            sampler = torch.utils.data.sampler.SequentialSampler(dataset)
        elif self.use_distributed:
            sampler = torch.utils.data.distributed.DistributedSampler(
                dataset, num_replicas=hvd.size(), rank=hvd.rank()
            )
        else:
            sampler = torch.utils.data.RandomSampler(dataset)

        data_loader = torch.utils.data.DataLoader(
            dataset, batch_size=batch_size, sampler=sampler, **kwargs
        )

        return data_loader

        Method to create a data loader for a given Tensor dataset.

        Args:
            mode(str): Mode for creating data loader. Could be train or test.
            dataset(torch.utils.data.Dataset): A Tensor dataset.
            batch_size(int): Batch size.

        Returns:
            torch.utils.data.DataLoader: A torch data loader to the given dataset.

        
utils_nlp/models/bert/sequence_classification_distributed.py