BatchEncoding Sample Input Error

[BrandonGroth]

When preparing a data sample:

def input_tiny():

# Random tokens and attention mask == 1
random_tokens = torch.randint(low=0, high=vocab_size, size=size)
attention_mask = torch.ones(size)

return BatchEncoding(
        {
            "input_ids": random_tokens,
            "attention_mask": attention_mask,
        }
    )

I get the following error in prepare_data_4_fwd() when invoking qmodel_prep():

E               TypeError: Data/dataloader provided by the user seems to have a type not in (tuple, list, or Tensor). 
E               If the data fetched from dataloader needs extra processing before being fed to model. Please consider using a `prefwdproc` function such that 
E                      model( prefwdproc(data_fetched_from_loader) )

BatchEncoding already has a lot of checks to "unpack it", so this is likely a bug.