enerzyme.models.nequip.core.NequIPWrapper

class enerzyme.models.nequip.core.NequIPWrapper(default_dtype: str, model_dtype: str, r_max: float, model_builders: List[str], num_layers: int, l_max: int, parity: bool, num_features: int, nonlinearity_type: str, resnet: bool, activation: str, nonlinearity_scalars: Dict[str, str], nonlinearity_gates: Dict[str, str], num_basis: int, BesselBasis_trainable: bool, PolynomialCutoff_p: int, invariant_layers: int, invariant_neurons: int, avg_num_neighbors: float, use_sc: bool, chemical_symbols: List[str])

Bases: BaseFFCore

__init__(default_dtype: str, model_dtype: str, r_max: float, model_builders: List[str], num_layers: int, l_max: int, parity: bool, num_features: int, nonlinearity_type: str, resnet: bool, activation: str, nonlinearity_scalars: Dict[str, str], nonlinearity_gates: Dict[str, str], num_basis: int, BesselBasis_trainable: bool, PolynomialCutoff_p: int, invariant_layers: int, invariant_neurons: int, avg_num_neighbors: float, use_sc: bool, chemical_symbols: List[str])

build(built_layers) → None

get_output(Ra: Tensor, Za: Tensor, batch_seg: Tensor) → Dict[str, Tensor]