From 58ad56afece3ea7faec2f1b7f68d90195f316d78 Mon Sep 17 00:00:00 2001 From: David de la Iglesia Castro Date: Mon, 15 Apr 2024 16:16:17 +0200 Subject: [PATCH] Use `step` interval in `estimated_stepping_batches` docs example (#19774) --- docs/source-pytorch/common/trainer.rst | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/docs/source-pytorch/common/trainer.rst b/docs/source-pytorch/common/trainer.rst index 0ad4592754134..c40ad1fcf92e2 100644 --- a/docs/source-pytorch/common/trainer.rst +++ b/docs/source-pytorch/common/trainer.rst @@ -1229,7 +1229,10 @@ dataloader if hadn't been set up already. optimizer = ... stepping_batches = self.trainer.estimated_stepping_batches scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, max_lr=1e-3, total_steps=stepping_batches) - return [optimizer], [scheduler] + return { + "optimizer": optimizer, + "lr_scheduler": {"scheduler": scheduler, "interval": "step"}, + } state *****