scvi.model.base.PyroJitGuideWarmup.on_before_zero_grad

PyroJitGuideWarmup.on_before_zero_grad(trainer, pl_module, optimizer)

Called after optimizer.step() and before optimizer.zero_grad().

Return type

NoneNone