Skip to content

Commit

Permalink
switch to eval mode to avoid create_graph in 'fit_output_to_model_out…
Browse files Browse the repository at this point in the history
…put'
  • Loading branch information
HydrogenSulfate committed Dec 22, 2024
1 parent 61e05e2 commit 79789da
Showing 1 changed file with 11 additions and 0 deletions.
11 changes: 11 additions & 0 deletions deepmd/pd/infer/deep_eval.py
Original file line number Diff line number Diff line change
Expand Up @@ -354,6 +354,12 @@ def _eval_model(
request_defs: list[OutputVariableDef],
):
model = self.dp.to(DEVICE)

# convert to eval mode
is_training = model.training
if is_training:
model.eval()

prec = NP_PRECISION_DICT[RESERVED_PRECISON_DICT[GLOBAL_PD_FLOAT_PRECISION]]

nframes = coords.shape[0]
Expand Down Expand Up @@ -419,6 +425,11 @@ def _eval_model(
results.append(
np.full(np.abs(shape), np.nan, dtype=prec)
) # this is kinda hacky

# restore to training mode if previously enabled
if is_training:
model.train()

return tuple(results)

def _eval_model_spin(
Expand Down

0 comments on commit 79789da

Please sign in to comment.