Simplified the model by always computing batch-first

modeling_norbert.py

@@ -156,8 +156,8 @@ class Attention(nn.Module):
         value = self.in_proj_v(hidden_states)  # shape: [B, T, D]
 
         # Reshape to [B, num_heads, T, head_size]
-        query = query.reshape(batch_size, query_len, self.num_heads, self.head_size).transpose(1, 2)  # shape: [B, num_heads, T_q, head_size]
-        key = key.reshape(batch_size, key_len, self.num_heads, self.head_size).permute(0, 2, 3, 1)  # shape: [B, num_heads, head_size, T_k]
+        query = query.view(batch_size, query_len, self.num_heads, self.head_size).transpose(1, 2)  # shape: [B, num_heads, T_q, head_size]
+        key = key.view(batch_size, key_len, self.num_heads, self.head_size).permute(0, 2, 3, 1)  # shape: [B, num_heads, head_size, T_k]
         value = value.view(batch_size, key_len, self.num_heads, self.head_size).transpose(1, 2)  # shape: [B, num_heads, T_k, head_size]
 
         # Compute relative positional contributions