[BUG] Warn users if they have too many GPUs

[alexbarghi-nv]

We have an ongoing issue on very small datasets where there are too few batches to evenly distribute the work across GPUs. We should automatically detect this case and show a warning instead of hanging.