code: def __init__(self):
      super().__init__()
      self.generator = Generator()
      self.discriminator = Discriminator()
      self.g_optimizer = optax.adam(2e-4, b1=0.5)
      self.d_optimizer = optax.adam(2e-4, b1=0.5)
  
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 1 to create a new WGAN_GP object from inherited base classes: elegy with example: obj = WGAN_GP()