chore: tweak step_lr param in e2e tests
@@ -32,7 +32,7 @@ class OneVarLM(pl.LightningModule):
def configure_optimizers(self):
opt = torch.optim.SGD(self.model.parameters(), self.lr)
- sched = torch.optim.lr_scheduler.StepLR(opt, step_size=1e-4)
+ sched = torch.optim.lr_scheduler.StepLR(opt, step_size=1, gamma=1e-6)
return [opt], [sched]
def training_step(self, batch, batch_idx, *args, **kwargs):
