Interestingly, it seems you are getting the error, if the batches are imbalanced, i.e. if the data cannot be equally split to the devices.
I just tested my code with a batch size of 5 (which is your last case before the error is raised) and it still works fine. Could you do the same and check, if you are seeing an error?
Also, are you using the latest PyTorch version?