BATCH_SIZE = 32

def batched_sample_data(inputs, labels):
    sampled_inputs = []
    sampled_labels = []
    for i in range(inputs.size(0)):
        inp, lab = sample_data(inputs[i], labels[i])
        sampled_inputs.append(inp)
        sampled_labels.append(lab)
    return torch.stack(sampled_inputs), torch.stack(sampled_labels)