# Buggy code
def auto_squeeze_dim_zeros(output):
    """
    In DP or DDP2 we need to unsqueeze dim 0
    :param output:
    :return:
    """
    for k, v in output.items():
        if not isinstance(v, torch.Tensor):
            continue

        is_scalar = v.dim() == 0
        if is_scalar:
            output[k] = output[k].unsqueeze(0)


# Traceback (RuntimeError)
Traceback (most recent call last):
File "./_debug_pytorch_lightning_16gpu.py", line 28, in <module>
main()
File "./_debug_pytorch_lightning_16gpu.py", line 24, in main
trainer.fit(model)
File "/home/ubuntu/anaconda3/envs/trfm/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 973, in fit
self.dp_train(model)
File "/home/ubuntu/anaconda3/envs/trfm/lib/python3.7/site-packages/pytorch_lightning/trainer/distrib_parts.py", line 267, in dp_train
self.run_pretrain_routine(model)
File "/home/ubuntu/anaconda3/envs/trfm/lib/python3.7/site-packages/pytorch_lightning/trainer/trainer.py", line 1139, in run_pretrain_routine
False)
File "/home/ubuntu/anaconda3/envs/trfm/lib/python3.7/site-packages/pytorch_lightning/trainer/evaluation_loop.py", line 291, in _evaluate
output = self.evaluation_forward(model, batch, batch_idx, dataloader_idx, test_mode)
File "/home/ubuntu/anaconda3/envs/trfm/lib/python3.7/site-packages/pytorch_lightning/trainer/evaluation_loop.py", line 441, in evaluation_forward
output = model(*args)
File "/home/ubuntu/anaconda3/envs/trfm/lib/python3.7/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
result = self.forward(*input, **kwargs)
File "/home/ubuntu/anaconda3/envs/trfm/lib/python3.7/site-packages/pytorch_lightning/overrides/data_parallel.py", line 65, in forward
outputs = self.parallel_apply(replicas, inputs, kwargs)
File "/home/ubuntu/anaconda3/envs/trfm/lib/python3.7/site-packages/pytorch_lightning/overrides/data_parallel.py", line 69, in parallel_apply
return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])
File "/home/ubuntu/anaconda3/envs/trfm/lib/python3.7/site-packages/pytorch_lightning/overrides/data_parallel.py", line 209, in parallel_apply
raise output
File "/home/ubuntu/anaconda3/envs/trfm/lib/python3.7/site-packages/pytorch_lightning/overrides/data_parallel.py", line 172, in _worker
output = module.validation_step(*input, **kwargs)
File "/home/ubuntu/anaconda3/envs/trfm/lib/python3.7/site-packages/pl_examples/models/lightning_template.py", line 102, in validation_step
val_loss = F.cross_entropy(y_hat, y)
File "/home/ubuntu/anaconda3/envs/trfm/lib/python3.7/site-packages/torch/nn/functional.py", line 2422, in cross_entropy
return nll_loss(log_softmax(input, 1), target, weight, None, ignore_index, None, reduction)
File "/home/ubuntu/anaconda3/envs/trfm/lib/python3.7/site-packages/torch/nn/functional.py", line 2218, in nll_loss
ret = torch._C._nn.nll_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index)
RuntimeError: Assertion `THCTensor_(checkGPU)(state, 4, input, target, output, total_weight)' failed. Some of weight/gradient/input tensors are located on different GPUs. Please move them to a single one. at /opt/conda/conda-bld/pytorch_1593673617738/work/aten/src/THCUNN/generic/ClassNLLCriterion.cu:31

# Fixed code
def auto_squeeze_dim_zeros(output):
    """
    In DP or DDP2 we need to unsqueeze dim 0
    :param output:
    :return:
    """
    if isinstance(output, torch.Tensor):
        output = output.unsqueeze(0)
        return output

    for k, v in output.items():
        if not isinstance(v, torch.Tensor):
            continue

        is_scalar = v.dim() == 0
        if is_scalar:
            output[k] = output[k].unsqueeze(0)