# Buggy code
def forward(self, *inputs, **kwargs):  # pragma: no-cover
    self._sync_params()
    if self.device_ids:
        inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)
        if len(self.device_ids) == 1:
            # --------------
            # LIGHTNING MOD
            # --------------
            # normal
            # output = self.module(*inputs[0], **kwargs[0])
            # lightning
            if self.module.training:
                output = self.module.training_step(*inputs[0], **kwargs[0])
            elif self.module.testing:
                output = self.module.test_step(*inputs[0], **kwargs[0])
            else:
                output = self.module.validation_step(*inputs[0], **kwargs[0])
        else:
            outputs = self.parallel_apply(
                self._module_copies[: len(inputs)], inputs, kwargs
            )
            output = self.gather(outputs, self.output_device)
    else:
        # output = self.module(*inputs, **kwargs)
        # normal lightning (ddp_cpu)
        if self.module.training:
            output = self.module.training_step(*inputs, **kwargs)
        elif self.module.testing:
            output = self.module.test_step(*inputs, **kwargs)
        else:
            output = self.module.validation_step(*inputs, **kwargs)

    if torch.is_grad_enabled():
        # We'll return the output object verbatim since it is a freeform
        # object. We need to find any tensors in this object, though,
        # because we need to figure out which parameters were used during
        # this forward pass, to ensure we short circuit reduction for any
        # unused parameters. Only if `find_unused_parameters` is set.
        if self.find_unused_parameters:
            self.reducer.prepare_for_backward(list(_find_tensors(output)))
        else:
            self.reducer.prepare_for_backward([])
    return output


# 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 forward(self, *inputs, **kwargs):  # pragma: no-cover
    self._sync_params()
    fx_called: str = ""

    if self.device_ids:
        inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)
        if len(self.device_ids) == 1:
            # --------------
            # LIGHTNING MOD
            # --------------
            # normal
            # output = self.module(*inputs[0], **kwargs[0])
            # lightning
            if self.module.training:
                output = self.module.training_step(*inputs[0], **kwargs[0])
                fx_called = "training_step"
            elif self.module.testing:
                output = self.module.test_step(*inputs[0], **kwargs[0])
                fx_called = "test_step"
            else:
                output = self.module.validation_step(*inputs[0], **kwargs[0])
                fx_called = "validation_step"
        else:
            outputs = self.parallel_apply(
                self._module_copies[: len(inputs)], inputs, kwargs
            )
            output = self.gather(outputs, self.output_device)
    else:
        # output = self.module(*inputs, **kwargs)
        # normal lightning (ddp_cpu)
        if self.module.training:
            output = self.module.training_step(*inputs, **kwargs)
        elif self.module.testing:
            output = self.module.test_step(*inputs, **kwargs)
        else:
            output = self.module.validation_step(*inputs, **kwargs)

    if torch.is_grad_enabled():
        # We'll return the output object verbatim since it is a freeform
        # object. We need to find any tensors in this object, though,
        # because we need to figure out which parameters were used during
        # this forward pass, to ensure we short circuit reduction for any
        # unused parameters. Only if `find_unused_parameters` is set.
        if self.find_unused_parameters:
            self.reducer.prepare_for_backward(list(_find_tensors(output)))
        else:
            self.reducer.prepare_for_backward([])

    if output is None:
        warn_missing_output(fx_called)

        m = f"{fx_called} returned None. Did you forget to re"
    return output