def METHOD_NAME(self):
    sizes_to_test = [
        # ((1,),),
        ((3, 2, 10, 3), (1,)),
        # ((20,),),
    ]
    list_of_classes = [
        list(range(100)),
        ["%d" % i for i in range(100)],
        list(range(300, 400)),
    ]
    for batch_size in [None, 1, 4]:
        for sizes, classes in zip(sizes_to_test, list_of_classes):
            def create_batch_labels(batch_idx):
                if batch_size is None:
                    # batch_size = 1
                    return classes[batch_idx]
                return classes[
                    batch_idx * batch_size : (batch_idx + 1) * batch_size
                ]
            bs = 1 if batch_size is None else batch_size
            num_batches = len(classes) // bs
            sizes_plus_batch = tuple((bs,) + si for si in sizes)
            data = [
                (
                    tuple(torch.randn(si) for si in sizes_plus_batch),
                    create_batch_labels(batch_idx),
                )
                for batch_idx in range(num_batches)
            ]
            with self.subTest(
                batch_size=batch_size, sizes=sizes_plus_batch, classes=classes
            ):
                self.class_test(data, classes, sizes)