diff --git a/deepctr_torch/layers/interaction.py b/deepctr_torch/layers/interaction.py index 21f27ef..74ce3f7 100644 --- a/deepctr_torch/layers/interaction.py +++ b/deepctr_torch/layers/interaction.py @@ -372,8 +372,8 @@ def __init__(self, embedding_size, head_num=2, use_res=True, scaling=False, if self.use_res: self.W_Res = nn.Parameter(torch.Tensor(embedding_size, output_size)) for tensor in self.parameters(): - nn.init.trunc_normal_(tensor, mean=0.0, std=0.05, - a=-0.1, b=0.1) + nn.init.trunc_normal_(tensor, mean=0.0, std=1.0, + a=-2.0, b=2.0) self.to(device) diff --git a/tests/layers/interaction_test.py b/tests/layers/interaction_test.py index ef43b04..321c212 100644 --- a/tests/layers/interaction_test.py +++ b/tests/layers/interaction_test.py @@ -76,7 +76,8 @@ def test_other_interaction_layers_match_tf(): layer = InteractingLayer(embedding_size=8, head_num=2, use_res=True) for parameter in layer.parameters(): - assert torch.max(torch.abs(parameter)) <= 0.1 + assert torch.max(torch.abs(parameter)) <= 2.0 + assert parameter.std() > 0.1 layer = InteractingLayer( embedding_size=6, head_num=2, att_embedding_size=4, use_res=True)