Unnecessarily strict requirement for batch size in MASPlugin

[niniack]

Describe the bug

In the MASPlugin, the following check seems to be unnecessarily strict:

Link to code

for _, batch in enumerate(dataloader):
    # Get batch
    if len(batch) == 2 or len(batch) == 3:
        x, _, t = batch[0], batch[1], batch[-1]
    else:
        raise ValueError("Batch size is not valid")

As a result, if I have a custom dataset which returns additional elements (two images and a label), thi

Expected behavior

I am not sure MAS needs to follow this check, it can simply grab the needed information (image, label, and token) and continue onwards.