def pre_backward(self, closure_loss: torch.Tensor) -> None:
"""Run before precision plugin executes backward."""
if not self.lightning_module.automatic_optimization:
prepare_for_backward(self.model, closure_loss)
def pre_backward(self, closure_loss: torch.Tensor, should_accumulate: bool,
optimizer: Optimizer, opt_idx: int):
"""Run before precision plugin executes backward"""
if not self.lightning_module.automatic_optimization:
prepare_for_backward(self.model, closure_loss)
def on_before_manual_backward(self, model: DistributedDataParallel,
output: Any):
prepare_for_backward(model, output)
