enerzyme.models.schnet.core.SchNetCore

class enerzyme.models.schnet.core.SchNetCore(hidden_channels: int = 128, dim_embedding: int = 128, num_interactions: int = 4, num_rbf: int = 128, cutoff_sr: float = 5.0, activation_fn: Literal['shifted_softplus', 'swish'] = 'shifted_softplus', activation_params: Dict[Literal['dim_feature', 'initial_alpha', 'initial_beta', 'learnable'], int | float | bool] = {}, shallow_ensemble_size: int = 1)

Bases: BaseFFCore

__init__(hidden_channels: int = 128, dim_embedding: int = 128, num_interactions: int = 4, num_rbf: int = 128, cutoff_sr: float = 5.0, activation_fn: Literal['shifted_softplus', 'swish'] = 'shifted_softplus', activation_params: Dict[Literal['dim_feature', 'initial_alpha', 'initial_beta', 'learnable'], int | float | bool] = {}, shallow_ensemble_size: int = 1)

build(built_layers: List[Module]) → None

get_output(idx_i_sr: Tensor, idx_j_sr: Tensor, Dij_sr: Tensor, rbf: Tensor, atom_embedding: Tensor)

reset_parameters()

Resets all learnable parameters of the module.