ValueError: only one element tensors can be converted to Python scalars
🐛 Bug

This happens in the training loop.
ValueError: only one element tensors can be converted to Python scalars

To Reproduce
From my observation, I believe this happens when the batch size can't be divided by gpu num. For example on the last batch of each epoch, and when you have 4 gpus but set batch size to 2.

Additional context
I think it would be nice to use only some of the gpus the user specified, while printing out a msg telling them that the gpus are not specified correctly. Current implementaion simply throws a not friendly error