def _create_custom_dataloader(self):

    class MtUcertaintyIterator():
        'Sample tasks using uncertainty measure.'

        def __init__(self, my_loader):
            self.my_loader = my_loader
            self.loader_iters = [iter(loader) for loader in self.my_loader.loaders]
            self.loader_iter_sizes = [len(i) for i in self.loader_iters]
            self.max_count = len(self.my_loader)
            self.batch_count = 0

        def __iter__(self):
            return self

        def __next__(self):
            if (self.batch_count == self.max_count):
                self.batch_count = 0
                raise StopIteration()
            test_batch = {}
            for (idx, loader_iter) in enumerate(self.loader_iters):
                try:
                    batch = loader_iter.__next__()
                except StopIteration:
                    new_loader_iter = iter(self.my_loader.loaders[idx])
                    self.loader_iters[idx] = new_loader_iter
                    batch = new_loader_iter.__next__()
                test_batch = self.batchify_data(batch, test_batch)
            inputs = {}
            for (k, v) in test_batch.items():
                if (k not in ['labels']):
                    inputs[k] = v.detach().to(self.my_loader.args.device)
            with torch.no_grad():
                model.select_batch_mode = True
                outputs = model(**inputs)
                model.select_batch_mode = False
            (test_batch_entropy, test_batch_entropy_mean, max_mean_batch_entropy) = outputs[(- 3):]
            for (_, v) in inputs.items():
                del v
            del inputs
            test_batch_entropy_mean = (test_batch_entropy_mean / max_mean_batch_entropy)
            test_batch_entropy = (test_batch_entropy * test_batch_entropy_mean)
            select_size = min(self.my_loader.args.train_batch_size, test_batch['input_ids'].shape[0])
            top_entropy = torch.topk(test_batch_entropy, select_size)
            for (k, v) in test_batch.items():
                test_batch[k] = torch.index_select(v, 0, top_entropy.indices)
            self.batch_count += 1
            return test_batch

        @staticmethod
        def batchify_data(data, curr_batch):
            for k in data.keys():
                if (k in curr_batch.keys()):
                    curr_batch[k] = torch.cat((curr_batch[k], data[k]), dim=0)
                else:
                    curr_batch[k] = data[k]
            return curr_batch

    class CustomLoader():

        def __init__(self, loaders, datasets, loader_args):
            self.loaders = loaders
            self.dataset = datasets
            self.args = loader_args
            self.current_epoch = 0

        def __iter__(self):
            iterator = MtUcertaintyIterator(self)
            for l in self.loaders:
                if isinstance(l.sampler, DistributedSampler):
                    l.sampler.set_epoch(self.current_epoch)
            self.current_epoch += 1
            return iterator

        def __len__(self):
            loader_len = [len(loader) for loader in self.loaders]
            if self.args.uniform_mt_sampling:
                return int((((self.args.percent_of_max_data_size * max(loader_len)) * len(self.loaders)) / self.args.train_batch_size))
            elif self.args.uncert_batch:
                return int(((max(loader_len) * len(self.loaders)) * self.args.percent_of_max_data_size))
            else:
                return sum(loader_len)
    model = self.model
    tasks = self.data_args.tasks
    data_loaders = []
    for dataset in self.train_dataset.datasets:
        train_sampler = (RandomSampler(dataset) if (self.args.local_rank == (- 1)) else DistributedSampler(dataset))
        data_loader = DataLoader(dataset, batch_size=self.args.train_batch_size, sampler=train_sampler, collate_fn=self.data_collator.collate_batch)
        data_loaders.append(data_loader)
    return CustomLoader(data_loaders, self.train_dataset, self.args)
def _create_custom_dataloader(self):

    class MtUcertaintyIterator():
        'Sample tasks using uncertainty measure.'

        def __init__(self, my_loader):
            self.my_loader = my_loader
            self.loader_iters = [iter(loader) for loader in self.my_loader.loaders]
            self.loader_iter_sizes = [len(i) for i in self.loader_iters]
            self.max_count = len(self.my_loader)
            self.batch_count = 0

        def __iter__(self):
            return self

        def __next__(self):
            if (self.batch_count == self.max_count):
                self.batch_count = 0
                raise StopIteration()
            test_batch = {}
            for (idx, loader_iter) in enumerate(self.loader_iters):
                try:
                    batch = loader_iter.__next__()
                except StopIteration:
                    new_loader_iter = iter(self.my_loader.loaders[idx])
                    self.loader_iters[idx] = new_loader_iter
                    batch = new_loader_iter.__next__()
                test_batch = self.batchify_data(batch, test_batch)
            inputs = {}
            for (k, v) in test_batch.items():
                if (k not in ['labels']):
                    inputs[k] = v.detach().to(self.my_loader.args.device)
            with torch.no_grad():
                model.select_batch_mode = True
                outputs = model(**inputs)
                model.select_batch_mode = False
            (test_batch_entropy, test_batch_entropy_mean, max_mean_batch_entropy) = outputs[(- 3):]
            for (_, v) in inputs.items():
                del v
            del inputs
            test_batch_entropy_mean = (test_batch_entropy_mean / max_mean_batch_entropy)
            test_batch_entropy = (test_batch_entropy * test_batch_entropy_mean)
            select_size = min(self.my_loader.args.train_batch_size, test_batch['input_ids'].shape[0])
            top_entropy = torch.topk(test_batch_entropy, select_size)
            for (k, v) in test_batch.items():
                test_batch[k] = torch.index_select(v, 0, top_entropy.indices)
            self.batch_count += 1
            return test_batch

        @staticmethod
        def batchify_data(data, curr_batch):
            for k in data.keys():
                if (k in curr_batch.keys()):
                    curr_batch[k] = torch.cat((curr_batch[k], data[k]), dim=0)
                else:
                    curr_batch[k] = data[k]
            return curr_batch

    class CustomLoader():

        def __init__(self, loaders, datasets, loader_args):
            self.loaders = loaders
            self.dataset = datasets
            self.args = loader_args
            self.current_epoch = 0

        def __iter__(self):
            iterator = MtUcertaintyIterator(self)
            for l in self.loaders:
                if isinstance(l.sampler, DistributedSampler):
                    l.sampler.set_epoch(self.current_epoch)
            self.current_epoch += 1
            return iterator

        def __len__(self):
            loader_len = [len(loader) for loader in self.loaders]
            if self.args.uniform_mt_sampling:
                return int((((self.args.percent_of_max_data_size * max(loader_len)) * len(self.loaders)) / self.args.train_batch_size))
            elif self.args.uncert_batch:
                return int(((max(loader_len) * len(self.loaders)) * self.args.percent_of_max_data_size))
            else:
                return sum(loader_len)
    model = self.model
    tasks = self.data_args.tasks
    data_loaders = []
    for dataset in self.train_dataset.datasets:
        train_sampler = (RandomSampler(dataset) if (self.args.local_rank == (- 1)) else DistributedSampler(dataset))
        data_loader = DataLoader(dataset, batch_size=self.args.train_batch_size, sampler=train_sampler, collate_fn=self.data_collator.collate_batch)
        data_loaders.append(data_loader)
    return CustomLoader(data_loaders, self.train_dataset, self.args)