# Buggy code
def optimizer_closure(self, split_batch, batch_idx, opt_idx, optimizer, hiddens):
    """
    wrap the forward step in a closure so second order methods work
    """
    # ---------------------------
    # FORWARD
    # ---------------------------
    with self.profiler.profile("model_forward"):
        if self.use_amp and NATIVE_AMP_AVALAIBLE:
            with torch.cuda.amp.autocast():
                training_step_output = self.training_forward(
                    split_batch, batch_idx, opt_idx, hiddens
                )
        else:
            training_step_output = self.training_forward(
                split_batch, batch_idx, opt_idx, hiddens
            )

        # ----------------------------
        # PROCESS THE RESULT
        # ----------------------------
        # format and reduce outputs accordingly
        training_step_output_for_epoch_end = training_step_output
        training_step_output = self.process_output(training_step_output, train=True)

        # TODO: temporary part of structured results PR
        training_step_output = AttributeDict(
            batch_loss=training_step_output[0],
            pbar_on_batch_end=training_step_output[1],
            log_metrics=training_step_output[2],
            callback_metrics=training_step_output[3],
            hiddens=training_step_output[4],
        )

        # if the user decides to finally reduce things in epoch_end, save raw output without graphs
        training_step_output_for_epoch_end = recursive_detach(
            training_step_output_for_epoch_end
        )

    # accumulate loss
    # (if accumulate_grad_batches = 1 no effect)
    closure_loss = training_step_output.batch_loss / self.accumulate_grad_batches

    # the loss will get scaled for amp. avoid any modifications to it
    untouched_loss = closure_loss.detach().clone()

    # backward pass
    model_ref = self.get_model()
    with self.profiler.profile("model_backward"):
        # scale loss for 16 bit
        if self.precision == 16 and not self.on_tpu:
            closure_loss = model_ref.amp_scale_loss(closure_loss, optimizer, opt_idx)

            # enter amp context
            if not NATIVE_AMP_AVALAIBLE:
                context = closure_loss
                closure_loss = closure_loss.__enter__()

        # do backward pass
        model_ref.backward(self, closure_loss, optimizer, opt_idx)

        # exit amp context
        if self.precision == 16 and not NATIVE_AMP_AVALAIBLE:
            a, b, c = None, None, None
            error = context.__exit__(a, b, c)
            if error:
                rank_zero_warn(a, b, c)
                raise Exception("apex unscale error")

        # once backward has been applied, release graph
        closure_loss = closure_loss.detach()
        training_step_output.batch_loss = training_step_output.batch_loss.detach()

    if self.use_horovod:
        # Synchronize Horovod to ensure gradient manipulations (e.g., loss scaling) are valid
        optimizer.synchronize()

    # insert after step hook
    if self.is_function_implemented("on_after_backward"):
        model_ref = self.get_model()
        with self.profiler.profile("on_after_backward"):
            model_ref.on_after_backward()

    result = AttributeDict(
        loss=untouched_loss,
        training_step_output=training_step_output,
        training_step_output_for_epoch_end=training_step_output_for_epoch_end,
        hiddens=training_step_output.hiddens,
    )
    return result


# Traceback (ModuleNotFoundError)
GPU available: False, used: False
TPU available: True, using: 8 TPU cores
Traceback (most recent call last):
File "bert_ner/light/fp16_debug.py", line 16, in <module>
trainer = pl.Trainer(tpu_cores=8, precision=16)
File "/anaconda3/envs/torch-xla-1.5/lib/python3.6/site-packages/pytorch_lightning/trainer/trainer.py", line 607, in __init__
self.init_amp()
File "/anaconda3/envs/torch-xla-1.5/lib/python3.6/site-packages/pytorch_lightning/trainer/auto_mix_precision.py", line 27, in init_amp
"You set `use_amp=True` but do not have apex installed."
ModuleNotFoundError: You set `use_amp=True` but do not have apex installed.Install apex first using this guide and rerun with use_amp=True:https://github.com/NVIDIA/apex#linux his run will NOT use 16 bit precision

# Fixed code
def optimizer_closure(self, split_batch, batch_idx, opt_idx, optimizer, hiddens):
    """
    wrap the forward step in a closure so second order methods work
    """
    # ---------------------------
    # FORWARD
    # ---------------------------
    with self.profiler.profile("model_forward"):
        if self.use_amp and NATIVE_AMP_AVALAIBLE and not self.use_tpu:
            with torch.cuda.amp.autocast():
                training_step_output = self.training_forward(
                    split_batch, batch_idx, opt_idx, hiddens
                )
        else:
            training_step_output = self.training_forward(
                split_batch, batch_idx, opt_idx, hiddens
            )

        # ----------------------------
        # PROCESS THE RESULT
        # ----------------------------
        # format and reduce outputs accordingly
        training_step_output_for_epoch_end = training_step_output
        training_step_output = self.process_output(training_step_output, train=True)

        # TODO: temporary part of structured results PR
        training_step_output = AttributeDict(
            batch_loss=training_step_output[0],
            pbar_on_batch_end=training_step_output[1],
            log_metrics=training_step_output[2],
            callback_metrics=training_step_output[3],
            hiddens=training_step_output[4],
        )

        # if the user decides to finally reduce things in epoch_end, save raw output without graphs
        training_step_output_for_epoch_end = recursive_detach(
            training_step_output_for_epoch_end
        )

    # accumulate loss
    # (if accumulate_grad_batches = 1 no effect)
    closure_loss = training_step_output.batch_loss / self.accumulate_grad_batches

    # the loss will get scaled for amp. avoid any modifications to it
    untouched_loss = closure_loss.detach().clone()

    # backward pass
    model_ref = self.get_model()
    with self.profiler.profile("model_backward"):
        # scale loss for 16 bit
        if self.precision == 16 and not self.on_tpu:
            closure_loss = model_ref.amp_scale_loss(closure_loss, optimizer, opt_idx)

            # enter amp context
            if not NATIVE_AMP_AVALAIBLE:
                context = closure_loss
                closure_loss = closure_loss.__enter__()

        # do backward pass
        model_ref.backward(self, closure_loss, optimizer, opt_idx)

        # exit amp context
        if self.precision == 16 and not NATIVE_AMP_AVALAIBLE and not self.on_tpu:
            a, b, c = None, None, None
            error = context.__exit__(a, b, c)
            if error:
                rank_zero_warn(a, b, c)
                raise Exception("apex unscale error")

        # once backward has been applied, release graph
        closure_loss = closure_loss.detach()
        training_step_output.batch_loss = training_step_output.batch_loss.detach()

    if self.use_horovod:
        # Synchronize Horovod to ensure gradient manipulations (e.g., loss scaling) are valid
        optimizer.synchronize()

    # insert after step hook
    if self.is_function_implemented("on_after_backward"):
        model_ref = self.get_model()
        with self.profiler.profile("on_after_backward"):
            model_ref.on_after_backward()

    result = AttributeDict(
        loss=untouched_loss,
        training_step_output=training_step_output,
        training_step_output_for_epoch_end=training_step_output_for_epoch_end,
        hiddens=training_step_output.hiddens,
    )
    return result