Upload roberta_layers.py

roberta_layers.py

@@ -199,9 +199,9 @@ class RobertaSelfAttention(nn.Module):
         query_layer = self.transpose_for_scores(mixed_query_layer)
 
         if xformers_available:
-            query_layer = query_layer.transpose(1, 2)
-            key_layer = key_layer.transpose(1, 2)
-            value_layer = value_layer.transpose(1, 2)
+            query_layer = mixed_query_layer
+            key_layer = self.key(hidden_states)
+            value_layer = self.value(hidden_states)
             context_layer = xops.memory_efficient_attention(
                 query_layer, key_layer, value_layer,  p=self.dropout_prob
             )