revertme: debug shapes

Author: NeoLegends

i6_models/parts/conformer/mhsa_rel_pos.py

@@ -216,6 +216,10 @@ def forward(self, input_tensor: torch.Tensor, sequence_mask: torch.Tensor) -> to
         v = value_seq.view(batch_dim_size, -1, self.num_heads, self.embed_dim_per_head)  # [B, T, H, F']
 
         # attention matrix a and c is computed inside torch's sdpa
+        print("q_with_bias_u", q_with_bias_u.transpose(-3, -2))
+        print("k", k.transpose(-3, -2))
+        print("v", v.transpose(-3, -2))
+        print("attn_bd_mask_scaled", attn_bd_mask_scaled.transpose(-3, -2))
         attn_output = F.scaled_dot_product_attention(
             q_with_bias_u.transpose(-3, -2),  # [B, #heads, T, F']
             k.transpose(-3, -2),  # [B, #heads, T', F']