Upload roberta_layers.py

roberta_layers.py

@@ -29,7 +29,7 @@ from transformers.utils import (
 from transformers import RobertaConfig
 
 try:
-    import xformers
+    from xformers import ops as xops
     xformers_available=True
 except ImportError as e:
     xformers_available=False
@@ -199,7 +199,7 @@ class RobertaSelfAttention(nn.Module):
         query_layer = self.transpose_for_scores(mixed_query_layer)
 
         if xformers_available:
-            context_layer = xformers.ops.memory_efficient_attention(
+            context_layer = xops.memory_efficient_attention(
                 query_layer, key_layer, value_layer,  p=self.dropout_prob
             )
         else: