Cannot find parameter - Githubissues

Current code on Tpu V5 with Google jax[tpu] pip package results in a "cannot find parameter" issue:

Traceback (most recent call last):
  File "/home/sven/test/sdxl_tpu_lcm.py", line 107, in <module>
    generate(default_prompt, default_neg_prompt)
  File "/home/sven/test/sdxl_tpu_lcm.py", line 91, in generate
    images = pipeline(
  File "/home/sven/.local/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/pipeline_flax_stable_diffusion_xl.py", line 116, in __call__
    images = _p_generate(
  File "/home/sven/.local/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/pipeline_flax_stable_diffusion_xl.py", line 298, in _p_generate
    return pipe._generate(
  File "/home/sven/.local/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/pipeline_flax_stable_diffusion_xl.py", line 265, in _generate
    latents, _ = jax.lax.fori_loop(0, num_inference_steps, loop_body, (latents, scheduler_state))
  File "/home/sven/.local/lib/python3.10/site-packages/diffusers/pipelines/stable_diffusion_xl/pipeline_flax_stable_diffusion_xl.py", line 245, in loop_body
    noise_pred = self.unet.apply(
  File "/home/sven/.local/lib/python3.10/site-packages/diffusers/models/unet_2d_condition_flax.py", line 400, in __call__
    sample, res_samples = down_block(sample, t_emb, deterministic=not train)
  File "/home/sven/.local/lib/python3.10/site-packages/diffusers/models/unet_2d_blocks_flax.py", line 159, in __call__
    hidden_states = resnet(hidden_states, temb, deterministic=deterministic)
  File "/home/sven/.local/lib/python3.10/site-packages/diffusers/models/resnet_flax.py", line 110, in __call__
    hidden_states = self.conv1(hidden_states)
  File "/home/sven/.local/lib/python3.10/site-packages/flax/linen/linear.py", line 637, in __call__
    kernel = self.param(
flax.errors.ScopeParamNotFoundError: Could not find parameter named "kernel" in scope "/down_blocks_0/resnets_0/conv1". (https://flax.readthedocs.io/en/latest/api_reference/flax.errors.html#flax.errors.ScopeParamNotFoundError)```
rafaelsf80 / genai-sdxl-latent-consistency-lora-tpu

Cannot find parameter #2
