include ReGLU in TabTransformer

AnFreTh · AnFreTh · commit fcb17c11fd1e · 2024-07-15T10:45:58.000Z
diff --git a/mambular/base_models/tabtransformer.py b/mambular/base_models/tabtransformer.py
@@ -11,6 +11,7 @@
 )
 from ..configs.tabtransformer_config import DefaultTabTransformerConfig
 from .basemodel import BaseModel
+from ..arch_utils.transformer_utils import CustomTransformerEncoderLayer
 
 
 class TabTransformer(BaseModel):
@@ -91,7 +92,7 @@ def __init__(
             "num_embedding_activation", config.num_embedding_activation
         )
 
-        encoder_layer = nn.TransformerEncoderLayer(
+        encoder_layer = CustomTransformerEncoderLayer(
             d_model=self.hparams.get("d_model", config.d_model),
             nhead=self.hparams.get("n_heads", config.n_heads),
             batch_first=True,

Original file line number	Diff line number	Diff line change
`@@ -11,6 +11,7 @@`
`11`	`11`	`)`
`12`	`12`	`from ..configs.tabtransformer_config import DefaultTabTransformerConfig`
`13`	`13`	`from .basemodel import BaseModel`
	`14`	`+from ..arch_utils.transformer_utils import CustomTransformerEncoderLayer`
`14`	`15`
`15`	`16`
`16`	`17`	`class TabTransformer(BaseModel):`
`@@ -91,7 +92,7 @@ def __init__(`
`91`	`92`	`"num_embedding_activation", config.num_embedding_activation`
`92`	`93`	`)`
`93`	`94`
`94`		`- encoder_layer = nn.TransformerEncoderLayer(`
	`95`	`+ encoder_layer = CustomTransformerEncoderLayer(`
`95`	`96`	`d_model=self.hparams.get("d_model", config.d_model),`
`96`	`97`	`nhead=self.hparams.get("n_heads", config.n_heads),`
`97`	`98`	`batch_first=True,`