src/transformers/trainer_pt_utils.py/smp_forward_backward
    def smp_forward_backward(model, inputs, gradient_accumulation_steps=1, scaler=None):
        with torch.cuda.amp.autocast(enabled=(scaler is not None)):
            outputs = model(**inputs)

        loss = outputs["loss"] if isinstance(outputs, dict) else outputs[0]
        loss /= gradient_accumulation_steps
        if scaler is not None:
            loss = scaler.scale(loss).squeeze()

        model.backward(loss)
        return loss