Update xlayers

2021-05-22 23:04:24 +08:00
parent 5b09f059fd
commit 8109ed166a
6 changed files with 104 additions and 33 deletions
--- a/xautodl/xlayers/super_transformer.py
+++ b/xautodl/xlayers/super_transformer.py
@@ -45,6 +45,7 @@ class SuperTransformerEncoderLayer(SuperModule):
        norm_affine: bool = True,
        act_layer: Callable[[], nn.Module] = nn.GELU,
        order: LayerOrder = LayerOrder.PreNorm,
+        use_mask: bool = False,
    ):
        super(SuperTransformerEncoderLayer, self).__init__()
        mha = SuperSelfAttention(
@@ -54,6 +55,7 @@ class SuperTransformerEncoderLayer(SuperModule):
            qkv_bias=qkv_bias,
            attn_drop=drop,
            proj_drop=drop,
+            use_mask=use_mask,
        )
        mlp = SuperMLPv2(
            d_model,