Pass use_triton flag to module (#336)

facebookresearch · Jun 21, 2022 · 22c1bab · 22c1bab
1 parent 10d85bf
commit 22c1bab
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/xformers/factory/block_factory.py b/xformers/factory/block_factory.py
@@ -147,7 +147,9 @@ def __init__(self, config: xFormerEncoderConfig, **kwargs):
             config.layer_norm_style == LayerNormStyle.Pre
             and config.layer_position.is_last()
         ):
-            self.wrap_ff = PostNorm(config.dim_model, self.wrap_ff)
+            self.wrap_ff = PostNorm(
+                config.dim_model, self.wrap_ff, use_triton=config.use_triton
+            )
 
         # Simplicial embeddings are only used if specified, and on the last layer
         self.simplicial_embedding: Optional[SimplicialEmbedding] = None