Improve the coverage in MLP_norm class (#278)

* improve TensorNet model coverage * Update pyproject.toml Signed-off-by: Tsz Wai Ko <[email protected]> * Improve the unit test for SO(3) equivarance in TensorNet class * improve SO3Net model class coverage and simplify TensorNet implementations * improve the coverage in MLP_norm class --------- Signed-off-by: Tsz Wai Ko <[email protected]>
materialsvirtuallab · Jun 24, 2024 · fc0cbf6 · fc0cbf6
1 parent e862cba
commit fc0cbf6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/layers/test_core_and_embedding.py b/tests/layers/test_core_and_embedding.py
@@ -48,7 +48,7 @@ def test_gated_mlp(self, x):
 
     @pytest.mark.parametrize("normalization", ["layer", "graph"])
     def test_mlp_norm(self, x, graph, normalization):
-        layer = MLP_norm(dims=[10, 3], normalization=normalization)
+        layer = MLP_norm(dims=[10, 3], normalization=normalization, normalize_hidden=True)
         out = layer(x, g=graph).double()
         assert [out.size()[0], out.size()[1]] == [4, 3]
         assert out.mean().item() == pytest.approx(0, abs=1e-6)