Starting to test distributed train, fix issue with batch_size reduce

Author: rwightman

train.py

@@ -1142,7 +1142,10 @@ def _backward(_loss):
 
             if args.distributed:
                 # scale gradient btw distributed ranks, each one can have different batch size
-                global_batch_size = utils.reduce_tensor(torch.tensor(batch_size, device=device), 1) # SUM
+                global_batch_size = utils.reduce_tensor(
+                    torch.tensor(batch_size, device=device, dtype=torch.float32),
+                    1 # SUM
+                )
                 dist_scale = args.world_size * batch_size / global_batch_size
             else:
                 dist_scale = None