core.models.gemnet_oc.layers.force_scaler

Copyright (c) Meta, Inc. and its affiliates. This source code is licensed under the MIT license found in the LICENSE file in the root directory of this source tree.

Classes

ForceScaler

Scales up the energy and then scales down the forces

Module Contents

class core.models.gemnet_oc.layers.force_scaler.ForceScaler(init_scale: float = 2.0**8, growth_factor: float = 2.0, backoff_factor: float = 0.5, growth_interval: int = 2000, max_force_iters: int = 50, enabled: bool = True)

Scales up the energy and then scales down the forces to prevent NaNs and infs in calculations using AMP. Inspired by torch.cuda.amp.GradScaler.

scale_factor

growth_factor

backoff_factor

growth_interval

max_force_iters

enabled

finite_force_results = 0

scale(energy)

unscale(forces)

calc_forces(energy, pos)

calc_forces_and_update(energy, pos)

update() → None