fix: decoder attention

Attention layers in decoder were not used at the correct step

Author: borisdayma

vqgan_jax/modeling_flax_vqgan.py

@@ -225,12 +225,12 @@ def setup(self):
                                dtype=self.dtype)
 
   def __call__(self, hidden_states, temb=None, deterministic: bool = True):
-    for res_block in self.block:
+    for i, res_block in enumerate(self.block):
       hidden_states = res_block(hidden_states,
                                 temb,
                                 deterministic=deterministic)
-      for attn_block in self.attn:
-        hidden_states = attn_block(hidden_states)
+      if self.attn:
+        hidden_states = self.attn[i](hidden_states)
 
     if self.upsample is not None:
       hidden_states = self.upsample(hidden_states)