Simplified the model by always computing batch-first

modeling_norbert.py

@@ -125,7 +125,7 @@ class Attention(nn.Module):
             - torch.arange(config.max_position_embeddings, dtype=torch.long).unsqueeze(0)
         position_indices = self.make_log_bucket_position(position_indices, config.position_bucket_size, config.max_position_embeddings)
         position_indices = config.position_bucket_size - 1 + position_indices
-        self.register_buffer("position_indices", position_indices, persistent=False)
+        self.register_buffer("position_indices", position_indices.contiguous(), persistent=False)
 
         self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
         self.scale = 1.0 / math.sqrt(3 * self.head_size)