test before you push

poets-ai · Feb 11, 2021 · 8d1fdb0 · 8d1fdb0
1 parent 267e6eb
commit 8d1fdb0
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/examples/WGAN-GP/model.py b/examples/WGAN-GP/model.py
@@ -80,11 +80,11 @@ def pred_step(self, x, states):
 
     def train_step(self, x, states):
         # training the discriminator on every iteration
-        d_loss, gp, states = self.discriminator_step_jit(x, states)
+        d_loss, gp, states = self.discriminator_step(x, states)
 
         step = states.step + 1
         no_update = lambda args: (0.0, args[1])
-        do_update = lambda args: self.generator_step_jit(len(args[0]), args[1])
+        do_update = lambda args: self.generator_step(len(args[0]), args[1])
         g_loss, states = jax.lax.cond(step % 5 == 0, do_update, no_update, (x, states))
 
         return {"d_loss": d_loss, "g_loss": g_loss, "gp": gp}, states.update(step=step)