# Buggy code
def _worker(i, module, input, kwargs, device=None):
    torch.set_grad_enabled(grad_enabled)
    if device is None:
        device = get_a_var(input).get_device()
    try:
        with torch.cuda.device(device):
            # this also avoids accidental slicing of `input` if it is a Tensor
            if not isinstance(input, (list, tuple)):
                input = (input,)

            module = module.to(device)

            # ---------------
            # CHANGE
            if module.training:
                output = module.training_step(*input, **kwargs)

            elif module.testing:
                output = module.test_step(*input, **kwargs)

            else:
                output = module.validation_step(*input, **kwargs)

            if module.use_dp or module.use_ddp2:
                auto_squeeze_dim_zeros(output)
            # ---------------

        with lock:
            results[i] = output
    except Exception as ex:
        with lock:
            results[i] = ex


# 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 _worker(i, module, input, kwargs, device=None):
    torch.set_grad_enabled(grad_enabled)
    fx_called: str = ""
    if device is None:
        device = get_a_var(input).get_device()
    try:
        with torch.cuda.device(device):
            # this also avoids accidental slicing of `input` if it is a Tensor
            if not isinstance(input, (list, tuple)):
                input = (input,)

            module = module.to(device)

            # ---------------
            # CHANGE
            if module.training:
                output = module.training_step(*input, **kwargs)
                fx_called = "training_step"
            elif module.testing:
                output = module.test_step(*input, **kwargs)
                fx_called = "test_step"
            else:
                output = module.validation_step(*input, **kwargs)
                fx_called = "validation_step"

            if output is None:
                warn_missing_output(fx_called)

            if output is not None and (module.use_dp or module.use_ddp2):
                auto_squeeze_dim_zeros(output)
            # ---------------

        with lock:
            results[i] = output
    except Exception as ex:
        with lock:
            results[i] = ex