deepmd.pt.model.model.ener_model

Module Contents

Classes

EnergyModel

A base class to implement common methods for all the Models.

Attributes

DPEnergyModel_

deepmd.pt.model.model.ener_model.DPEnergyModel_

class deepmd.pt.model.model.ener_model.EnergyModel(*args, **kwargs)

Bases: deepmd.pt.model.model.dp_model.DPModelCommon, DPEnergyModel_

A base class to implement common methods for all the Models.

model_type = 'ener'

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

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