Skip to content

Commit

Permalink
remove painn parameters when using hydra
Browse files Browse the repository at this point in the history
  • Loading branch information
misko committed Aug 6, 2024
1 parent 62899bd commit b62d968
Showing 1 changed file with 2 additions and 1 deletion.
3 changes: 2 additions & 1 deletion src/fairchem/core/models/painn/painn.py
Original file line number Diff line number Diff line change
Expand Up @@ -671,7 +671,7 @@ def forward(self, x, v):
class PaiNNEnergyHead(nn.Module, HeadInterface):
def __init__(self, backbone):
super().__init__()

backbone.out_energy = None
self.out_energy = nn.Sequential(
nn.Linear(backbone.hidden_channels, backbone.hidden_channels // 2),
ScaledSiLU(),
Expand All @@ -697,6 +697,7 @@ def __init__(self, backbone):
self.direct_forces = backbone.direct_forces

if self.direct_forces:
backbone.out_forces = None
self.out_forces = PaiNNOutput(backbone.hidden_channels)

def forward(
Expand Down

0 comments on commit b62d968

Please sign in to comment.