Merge pull request #7 from borisdayma/fix_do_init

fix: uses _do_init arg

Author: patil-suraj

vqgan_jax/modeling_flax_vqgan.py

@@ -589,14 +589,16 @@ def __init__(
       input_shape: Tuple = (1, 256, 256, 3),
       seed: int = 0,
       dtype: jnp.dtype = jnp.float32,
+      _do_init: bool = True,
       **kwargs,
   ):
     module = self.module_class(config=config, dtype=dtype, **kwargs)
     super().__init__(config,
                      module,
                      input_shape=input_shape,
                      seed=seed,
-                     dtype=dtype)
+                     dtype=dtype,
+                     _do_init=_do_init)
 
   def init_weights(self, rng: jax.random.PRNGKey,
                    input_shape: Tuple) -> FrozenDict: