def training_step(self, split_batch, batch_idx, opt_idx, hiddens):
    model_ref = self.trainer.lightning_module
    with self.trainer.profiler.profile('model_forward'):
        args = self.build_train_args(split_batch, batch_idx, opt_idx, hiddens)
        model_ref._current_fx_name = 'training_step'
        model_ref._results = Result()
        with self.trainer.profiler.profile('training_step'):
            training_step_output = self.trainer.accelerator.training_step(args)
            self.trainer.accelerator.post_training_step()
        self.trainer.logger_connector.cache_logged_metrics()
        self._check_training_step_output(training_step_output)
        training_step_output = self.trainer.call_hook('training_step_end', training_step_output)
        (training_step_output_for_epoch_end, training_step_output) = self._process_training_step_output(training_step_output, split_batch)
        if (training_step_output_for_epoch_end is None):
            return
    closure_loss = None
    untouched_loss = None
    if self.trainer.lightning_module.automatic_optimization:
        closure_loss = (training_step_output.minimize / self.trainer.accumulate_grad_batches)
        untouched_loss = closure_loss.detach().clone()
    result = AttributeDict(closure_loss=closure_loss, loss=untouched_loss, training_step_output=training_step_output, training_step_output_for_epoch_end=training_step_output_for_epoch_end)
    return result
def training_step(self, split_batch, batch_idx, opt_idx, hiddens):
    model_ref = self.trainer.lightning_module
    with self.trainer.profiler.profile('model_forward'):
        args = self.build_train_args(split_batch, batch_idx, opt_idx, hiddens)
        model_ref._current_fx_name = 'training_step'
        model_ref._results = Result()
        with self.trainer.profiler.profile('training_step'):
            training_step_output = self.trainer.accelerator.training_step(args)
            self.trainer.accelerator.post_training_step()
        self.trainer.logger_connector.cache_logged_metrics()
        self._check_training_step_output(training_step_output)
        training_step_output = self.trainer.call_hook('training_step_end', training_step_output)
        (training_step_output_for_epoch_end, training_step_output) = self._process_training_step_output(training_step_output, split_batch)
        if (training_step_output_for_epoch_end is None):
            return
    closure_loss = None
    untouched_loss = None
    if self.trainer.lightning_module.automatic_optimization:
        closure_loss = (training_step_output.minimize / self.trainer.accumulate_grad_batches)
        untouched_loss = closure_loss.detach().clone()
    result = AttributeDict(closure_loss=closure_loss, loss=untouched_loss, training_step_output=training_step_output, training_step_output_for_epoch_end=training_step_output_for_epoch_end)
    return result