Split Sizes Throws Error

[jstremme]

@MichalBrzozowski91, thanks for this project! Really great stuff.

I get the following error in main.py on a four GPU setup when attempting to fine-tune a BERT model:

With batch size 1:
RuntimeError: split_with_sizes expects split_sizes to sum exactly to 800 (input tensor's size at dimension 0), but got split_sizes=[16]

With batch size 4:
RuntimeError: split_with_sizes expects split_sizes to sum exactly to 2750 (input tensor's size at dimension 0), but got split_sizes=[25, 6, 8, 16]

I would expect number_of_chunks to be of variable size for each record in the batch, but no matter my batch size, I seem to get an error at preds_split = preds.split(number_of_chunks) in main.py.

Any idea what I might be missing?

[MichalBrzozowski91]

Does an example work on your setup?

[jstremme]

Does an example work on your setup?

Apologies for the late reply. Your example code works with my setup. It must be the customizations I was adding to your code that aren't working.

I'll close this issue, as there doesn't appear to be anything wrong with your code.