Fix JaxTrainingPlan to update global_step (#1791)

martinkim0 · adamgayoso · web-flow · commit 4537b787da2d · 2022-11-17T01:20:07.000Z
* Add fake optimizer step to JaxTrainingPlan

* Update scvi/train/_trainingplans.py

Co-authored-by: Adam Gayoso &lt;adamgayoso@users.noreply.github.com&gt;

Co-authored-by: Adam Gayoso &lt;adamgayoso@users.noreply.github.com&gt;
diff --git a/scvi/train/_trainingplans.py b/scvi/train/_trainingplans.py
@@ -1143,6 +1143,9 @@ def training_step(self, batch, batch_idx):
             prog_bar=True,
         )
         self.compute_and_log_metrics(loss_output, self.train_metrics, "train")
+        # Update the dummy optimizer to update the global step
+        _opt = self.optimizers()
+        _opt.step()
 
     @partial(jax.jit, static_argnums=(0,))
     def jit_validation_step(

Original file line number	Diff line number	Diff line change
`@@ -1143,6 +1143,9 @@ def training_step(self, batch, batch_idx):`
`1143`	`1143`	`prog_bar=True,`
`1144`	`1144`	`)`
`1145`	`1145`	`self.compute_and_log_metrics(loss_output, self.train_metrics, "train")`
	`1146`	`+ # Update the dummy optimizer to update the global step`
	`1147`	`+ _opt = self.optimizers()`
	`1148`	`+ _opt.step()`
`1146`	`1149`
`1147`	`1150`	`@partial(jax.jit, static_argnums=(0,))`
`1148`	`1151`	`def jit_validation_step(`