Repository: pytorch
Issue: Numpy.float64 vs native python float breaks DDP
Labels: triaged,module: ddp

## 🐛 Bug

So I was running some code to do cutmix. It runs fine on a single GPU, but everything blows when going to DDP. The error DDP was spitting out was this:

```
<cut off>
   return tuple(torch._C._scatter(tensor, devices, chunk_sizes, dim, streams))
RuntimeError: chunk expects at least a 1-dimensional tensor
```
The full stack trace is below. After much painful debugging, I isolated the issue to a single line of code:
```
# from https://github.com/clovaai/CutMix-PyTorch/blob/master/train.py#L228
def cutmix(x, y, beta):
    lam = np.random.beta(beta, beta)
    batch_size = x.size()[0]
    rand_index = torch.randperm(batch_size, device=x.device)
    target_a = y
    target_b = y[rand_index]
    bbx1, bby1, bbx2, bby2 = rand_bbox(x.size(), lam)
    x[:, :, bbx1:bbx2, bby1:bby2] = x[rand_index, :, bbx1:bbx2, bby1:bby2]
    # adjust lambda to exactly match pixel ratio
    lam =1 - ((bbx2 - bbx1) * (bby2 - bby1) / (x.size()[-1] * x.size()[-2])) # **********************************

    return x, target_a, target_b, lam
```
where the starred line is the issue. When that line is commented out, the error goes away entirely. The type of `lam` is originally a native python float, on the first line of the function. Later, it becomes a `numpy.float64` after doing the `1 - such and such`. I initially tried converting `lam` to a torch float tensor - this too raises the same error when running. **However, this line is not when the issue occurs**. The issue actually occurs later, when data is being collated I believe? Some point after cutmix has run but before my model receives a single data batch. 

The only place where `lam` gets used after the data transformation is for calculating the loss value: 
```
def mixup_criterion(criterion, pred, y_a, y_b, lam):
    return lam * criterion(pred, y_a) + (1 - lam) * criterion(pred, y_b)
```
But the program never reaches this point before erroring out, so I don't understand why this would even cause a problem at all.

## To Reproduce

Pass along a `numpy.float64` value from your dataloader it would seem, while using DDP.

Full stack trace:
```
Traceback (most recent call last):
  File "classifier.py", line 63, in <module>
    main(args, config)
  File "classifier.py", line 43, in main
    debug_model(model, data, config, args)
  File "/home/ndalton/code/critical_view/utils.py", line 62, in debug_model
    trainer.fit(model, datamodule=data)
  File "/home/ndalton/miniconda3/envs/cvos/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 514, in fit
    self.dispatch()
  File "/home/ndalton/miniconda3/envs/cvos/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 554, in dispatch
    self.accelerator.start_training(self)
  File "/home/ndalton/miniconda3/envs/cvos/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 74, in start_training
    self.training_type_plugin.start_training(trainer)
  File "/home/ndalton/miniconda3/envs/cvos/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py", line 111, in start_training
    self._results = trainer.run_train()
  File "/home/ndalton/miniconda3/envs/cvos/lib/python3.8/site-packages/pytorch_lightning/trainer/trainer.py", line 645, in run_train
    self.train_loop.run_training_epoch()
  File "/home/ndalton/miniconda3/envs/cvos/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py", line 493, in run_training_epoch
    batch_output = self.run_training_batch(batch, batch_idx, dataloader_idx)
  File "/home/ndalton/miniconda3/envs/cvos/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py", line 655, in run_training_batch
    self.optimizer_step(optimizer, opt_idx, batch_idx, train_step_and_backward_closure)
  File "/home/ndalton/miniconda3/envs/cvos/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py", line 426, in optimizer_step
    model_ref.optimizer_step(
  File "/home/ndalton/miniconda3/envs/cvos/lib/python3.8/site-packages/pytorch_lightning/core/lightning.py", line 1384, in optimizer_step
    optimizer.step(closure=optimizer_closure)
  File "/home/ndalton/miniconda3/envs/cvos/lib/python3.8/site-packages/pytorch_lightning/core/optimizer.py", line 214, in step
    self.__optimizer_step(*args, closure=closure, profiler_name=profiler_name, **kwargs)
  File "/home/ndalton/miniconda3/envs/cvos/lib/python3.8/site-packages/pytorch_lightning/core/optimizer.py", line 134, in __optimizer_step
    trainer.accelerator.optimizer_step(optimizer, self._optimizer_idx, lambda_closure=closure, **kwargs)
  File "/home/ndalton/miniconda3/envs/cvos/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 278, in optimizer_step
    self.run_optimizer_step(optimizer, opt_idx, lambda_closure, **kwargs)
  File "/home/ndalton/miniconda3/envs/cvos/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 283, in run_optimizer_step
    self.training_type_plugin.optimizer_step(optimizer, lambda_closure=lambda_closure, **kwargs)
  File "/home/ndalton/miniconda3/envs/cvos/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/training_type_plugin.py", line 160, in optimizer_step
    optimizer.step(closure=lambda_closure, **kwargs)
  File "/home/ndalton/miniconda3/envs/cvos/lib/python3.8/site-packages/torch/autograd/grad_mode.py", line 26, in decorate_context
    return func(*args, **kwargs)
  File "/home/ndalton/miniconda3/envs/cvos/lib/python3.8/site-packages/torch/optim/adam.py", line 66, in step
    loss = closure()
  File "/home/ndalton/miniconda3/envs/cvos/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py", line 649, in train_step_and_backward_closure
    result = self.training_step_and_backward(
  File "/home/ndalton/miniconda3/envs/cvos/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py", line 743, in training_step_and_backward
    result = self.training_step(split_batch, batch_idx, opt_idx, hiddens)
  File "/home/ndalton/miniconda3/envs/cvos/lib/python3.8/site-packages/pytorch_lightning/trainer/training_loop.py", line 293, in training_step
    training_step_output = self.trainer.accelerator.training_step(args)
  File "/home/ndalton/miniconda3/envs/cvos/lib/python3.8/site-packages/pytorch_lightning/accelerators/accelerator.py", line 157, in training_step
    return self.training_type_plugin.training_step(*args)
  File "/home/ndalton/miniconda3/envs/cvos/lib/python3.8/site-packages/pytorch_lightning/plugins/training_type/ddp.py", line 287, in training_step
    return self.model(*args, **kwargs)
  File "/home/ndalton/miniconda3/envs/cvos/lib/python3.8/site-packages/torch/nn/modules/module.py", line 727, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/ndalton/miniconda3/envs/cvos/lib/python3.8/site-packages/torch/nn/parallel/distributed.py", line 617, in forward
    inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)
  File "/home/ndalton/miniconda3/envs/cvos/lib/python3.8/site-packages/torch/nn/parallel/distributed.py", line 643, in scatter
    return scatter_kwargs(inputs, kwargs, device_ids, dim=self.dim)
  File "/home/ndalton/miniconda3/envs/cvos/lib/python3.8/site-packages/torch/nn/parallel/scatter_gather.py", line 36, in scatter_kwargs
    inputs = scatter(inputs, target_gpus, dim) if inputs else []
  File "/home/ndalton/miniconda3/envs/cvos/lib/python3.8/site-packages/torch/nn/parallel/scatter_gather.py", line 28, in scatter
    res = scatter_map(inputs)
  File "/home/ndalton/miniconda3/envs/cvos/lib/python3.8/site-packages/torch/nn/parallel/scatter_gather.py", line 15, in scatter_map
    return list(zip(*map(scatter_map, obj)))
  File "/home/ndalton/miniconda3/envs/cvos/lib/python3.8/site-packages/torch/nn/parallel/scatter_gather.py", line 19, in scatter_map
    return list(map(type(obj), zip(*map(scatter_map, obj.items()))))
  File "/home/ndalton/miniconda3/envs/cvos/lib/python3.8/site-packages/torch/nn/parallel/scatter_gather.py", line 15, in scatter_map
    return list(zip(*map(scatter_map, obj)))
  File "/home/ndalton/miniconda3/envs/cvos/lib/python3.8/site-packages/torch/nn/parallel/scatter_gather.py", line 13, in scatter_map
    return Scatter.apply(target_gpus, None, dim, obj)
  File "/home/ndalton/miniconda3/envs/cvos/lib/python3.8/site-packages/torch/nn/parallel/_functions.py", line 92, in forward
    outputs = comm.scatter(input, target_gpus, chunk_sizes, ctx.dim, streams)
  File "/home/ndalton/miniconda3/envs/cvos/lib/python3.8/site-packages/torch/nn/parallel/comm.py", line 186, in scatter
    return tuple(torch._C._scatter(tensor, devices, chunk_sizes, dim, streams))
RuntimeError: chunk expects at least a 1-dimensional tensor

```

## Expected behavior

The program to run without error. The fix is actually quite easy, I just had to change the starred line to: `    lam = float(1 - ((bbx2 - bbx1) * (bby2 - bby1) / (x.size()[-1] * x.size()[-2])))`. Basically just wrap it in `float(...)`. However, it took several hours to isolate the issue.

## Environment

Please copy and paste the output from our
[environment collection script](https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py)
(or fill out the checklist below manually).

You can get the script and run it with:
```
wget https://raw.githubusercontent.com/pytorch/pytorch/master/torch/utils/collect_env.py
# For security purposes, please check the contents of collect_env.py before running it.
python collect_env.py
```

PyTorch version: 1.7.1
Is debug build: False
CUDA used to build PyTorch: 11.0
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.2 LTS (x86_64)
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
Clang version: Could not collect
CMake version: version 3.16.3

Python version: 3.8 (64-bit runtime)
Is CUDA available: True
CUDA runtime version: 11.1.105
GPU models and configuration: 
GPU 0: Quadro RTX 8000
GPU 1: Quadro RTX 8000
GPU 2: Quadro RTX 8000
GPU 3: Quadro RTX 8000

Nvidia driver version: 460.39
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.19.2
[pip3] pytorch-lightning==1.2.3
[pip3] torch==1.7.1
[pip3] torchvision==0.8.2
[pip3] vit-pytorch==0.6.7
[conda] blas                      1.0                         mkl  
[conda] cudatoolkit               11.0.221             h6bb024c_0  
[conda] mkl                       2020.4             h726a3e6_304    conda-forge
[conda] mkl-service               2.3.0            py38h1e0a361_2    conda-forge
[conda] mkl_fft                   1.3.0            py38h5c078b8_1    conda-forge
[conda] mkl_random                1.2.0            py38hc5bc63f_1    conda-forge
[conda] numpy                     1.19.2           py38h54aff64_0  
[conda] numpy-base                1.19.2           py38hfa32c7d_0  
[conda] pytorch                   1.7.1           py3.8_cuda11.0.221_cudnn8.0.5_0    pytorch
[conda] pytorch-lightning         1.2.3                    pypi_0    pypi
[conda] torchvision               0.8.2                py38_cu110    pytorch
[conda] vit-pytorch               0.6.7                    pypi_0    pypi