vocoder debug

Author: prishajain1

src/maxdiffusion/models/ltx2/vocoder_ltx2.py

@@ -470,6 +470,7 @@ def __call__(self, hidden_states: Array, time_last: bool = False) -> Array:
       print(f"After resnets level {i} - shape: {hidden_states.shape}, min: {hidden_states.min()}, max: {hidden_states.max()}")
 
     hidden_states = self.act_out(hidden_states)
+    print(f"After act_out - shape: {hidden_states.shape}, min: {hidden_states.min()}, max: {hidden_states.max()}")
     print(f"conv_out kernel - min: {self.conv_out.kernel.value.min()}, max: {self.conv_out.kernel.value.max()}")
     hidden_states = self.conv_out(hidden_states)
     print(f"After conv_out - shape: {hidden_states.shape}, min: {hidden_states.min()}, max: {hidden_states.max()}")