code: def train_step(self, x, states):
      # training the discriminator on every iteration
      d_loss, gp, states = self.discriminator_step(x, states)
  
      # training the generator only every 5 iterations as recommended in the original WGAN paper
      step = states.step + 1
      no_update = lambda states: (0.0, states)
      do_update = lambda states: self.generator_step(len(x), states)
      g_loss, states = jax.lax.cond(
          step % 5 == 0, do_update, no_update, states
      )
  
      return {
          "d_loss": d_loss,
          "g_loss": g_loss,
          "gp": gp,
      }, states.update(step=step)
  
file_path: elegy/examples/need-fixing/WGAN-GP/model.py
imports: import jax
  import jax.numpy
  import optax
  import elegy
label_desc: Use image node_id 4 for calling the WGAN_GP obj's underlying member method code with example usage: obj.train_step(x, states) and returns: dict,  states