Construct the DataLoader. Pad data, stack label and lengths

def get_dataloader(train_dataset, train_data_lengths,
                   test_dataset, batch_size):
    """ Construct the DataLoader. Pad data, stack label and lengths"""
    bucket_num, bucket_ratio = 20, 0.2
    batchify_fn = gluonnlp.data.batchify.Tuple(
        gluonnlp.data.batchify.Pad(axis=0, ret_length=True),
        gluonnlp.data.batchify.Stack(dtype='float32'))
    batch_sampler = gluonnlp.data.sampler.FixedBucketSampler(
        train_data_lengths,
        batch_size=batch_size,
        num_buckets=bucket_num,
        ratio=bucket_ratio,
        shuffle=True)
    train_dataloader = gluon.data.DataLoader(
        dataset=train_dataset,
        batch_sampler=batch_sampler,
        batchify_fn=batchify_fn)
    test_dataloader = gluon.data.DataLoader(
        dataset=test_dataset,
        batch_size=batch_size,
        shuffle=False,
        batchify_fn=batchify_fn)
    return train_dataloader, test_dataloader