added gMLP_norm option if gMLP is used at read_out

CederGroupHub · Sep 21, 2023 · 195d48c · 195d48c
1 parent b59d301
commit 195d48c
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/chgnet/model/model.py b/chgnet/model/model.py
@@ -292,6 +292,7 @@ def __init__(
                     hidden_dim=mlp_hidden_dims,
                     output_dim=mlp_hidden_dims[-1],
                     dropout=mlp_dropout,
+                    norm=gMLP_norm,
                     activation=non_linearity,
                 ),
                 nn.Linear(in_features=mlp_hidden_dims[-1], out_features=1),