import torch
    from torch.utils.data import IterableDataset
    from transformers import (
        AutoModelForMaskedLM,
        AutoTokenizer,
        DataCollatorForLanguageModeling,
        Trainer,
        TrainingArguments,
    )

    data = [
        {
            "input_ids": torch.tensor([101, 2040, 2001, 1999, 14936, 102]),
            "token_type_ids": torch.tensor([0, 0, 0, 0, 0, 0]),
            "attention_mask": torch.tensor([1, 1, 1, 1, 1, 1]),
        },
        {
            "input_ids": torch.tensor([101, 2040, 102]),
            "token_type_ids": torch.tensor([0, 0, 0]),
            "attention_mask": torch.tensor([1, 1, 1]),
        },
        {
            "input_ids": torch.tensor([101, 2040, 2001, 1999]),
            "token_type_ids": torch.tensor([0, 0, 0, 0]),
            "attention_mask": torch.tensor([1, 1, 1, 1]),
        },
        {
            "input_ids": torch.tensor([101, 2040, 2001, 1999, 14936, 102]),
            "token_type_ids": torch.tensor([0, 0, 0, 0, 0, 0]),
            "attention_mask": torch.tensor([1, 1, 1, 1, 1, 1]),
        },
        {
            "input_ids": torch.tensor([101]),
            "token_type_ids": torch.tensor([00]),
            "attention_mask": torch.tensor([1]),
        },
        {
            "input_ids": torch.tensor([101]),
            "token_type_ids": torch.tensor([00]),
            "attention_mask": torch.tensor([1]),
        },
    ]

    class ExampleDataset(IterableDataset):
        def __init__(self, data):
            super().__init__()
            self.data = data * 20

        def __iter__(self):
            for x in self.data:
                yield x

        def __len__(self):
            return len(self.data)

    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
    model = AutoModelForMaskedLM.from_pretrained("bert-base-cased")
    train_args = TrainingArguments(
        output_dir="output",
        num_train_epochs=3,
        per_device_train_batch_size=2,
    )
    dc = DataCollatorForLanguageModeling(tokenizer=tokenizer)

    trainer = Trainer(
        train_dataset=ExampleDataset(data),
        model=model,
        args=train_args,
        data_collator=dc,
    )
    trainer.train()

Traceback (most recent call last):
  File "fm_model/data/scratch.py", line 242, in <module>
    trainer.train()
  File "/opt/conda/envs/fmmodel/lib/python3.8/site-packages/transformers/trainer.py", line 1556, in train
    return inner_training_loop(
  File "/opt/conda/envs/fmmodel/lib/python3.8/site-packages/transformers/trainer.py", line 1816, in _inner_training_loop
    for step, inputs in enumerate(epoch_iterator):
  File "/opt/conda/envs/fmmodel/lib/python3.8/site-packages/accelerate/data_loader.py", line 597, in __iter__
    next_batch, next_batch_info = self._fetch_batches(main_iterator)
  File "/opt/conda/envs/fmmodel/lib/python3.8/site-packages/accelerate/data_loader.py", line 528, in _fetch_batches
    batch = concatenate(batches, dim=0)
  File "/opt/conda/envs/fmmodel/lib/python3.8/site-packages/accelerate/utils/operations.py", line 496, in concatenate
    return type(data[0])({k: concatenate([d[k] for d in data], dim=dim) for k in data[0].keys()})
  File "/opt/conda/envs/fmmodel/lib/python3.8/site-packages/accelerate/utils/operations.py", line 496, in <dictcomp>
    return type(data[0])({k: concatenate([d[k] for d in data], dim=dim) for k in data[0].keys()})
  File "/opt/conda/envs/fmmodel/lib/python3.8/site-packages/accelerate/utils/operations.py", line 499, in concatenate
    return torch.cat(data, dim=dim)
RuntimeError: Sizes of tensors must match except in dimension 0. Expected size 1 but got size 6 for tensor number 1 in the list.