Upload roberta_layers.py

roberta_layers.py

@@ -199,6 +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)
             context_layer = xops.memory_efficient_attention(
                 query_layer, key_layer, value_layer,  p=self.dropout_prob
             )