deepmd.pt_expt.model.dipole_model

Attributes

DPDipoleModel_

Classes

DipoleModel

Common methods for DP models.

Module Contents

deepmd.pt_expt.model.dipole_model.DPDipoleModel_

class deepmd.pt_expt.model.dipole_model.DipoleModel(*args: Any, **kwargs: Any)

Bases: deepmd.dpmodel.model.dp_model.DPModelCommon, DPDipoleModel_

Common methods for DP models.

This class provides common functionality for DeepPot models, including neighbor selection updates and fitting network access.

forward(coord: torch.Tensor, atype: torch.Tensor, box: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, charge_spin: torch.Tensor | None = None) → dict[str, torch.Tensor]

Default forward delegates to call().

Subclasses (e.g. EnergyModel) override this with output translation.

forward_lower(extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, charge_spin: torch.Tensor | None = None) → dict[str, torch.Tensor]

translated_output_def() → dict[str, Any]

forward_lower_exportable(extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, mapping: torch.Tensor | None = None, fparam: torch.Tensor | None = None, aparam: torch.Tensor | None = None, do_atomic_virial: bool = False, charge_spin: torch.Tensor | None = None, **make_fx_kwargs: Any) → torch.nn.Module