deepmd.pt.optimizer.KFWrapper

Module Contents

Classes

KFOptimizerWrapper

class deepmd.pt.optimizer.KFWrapper.KFOptimizerWrapper(model: torch.nn.Module, optimizer: torch.optim.optimizer.Optimizer, atoms_selected: int, atoms_per_group: int, is_distributed: bool = False)

update_energy(inputs: dict, Etot_label: torch.Tensor, update_prefactor: float = 1) → None

update_force(inputs: dict, Force_label: torch.Tensor, update_prefactor: float = 1) → None

update_denoise_coord(inputs: dict, clean_coord: torch.Tensor, update_prefactor: float = 1, mask_loss_coord: bool = True, coord_mask: torch.Tensor = None) → None

__sample(atoms_selected: int, atoms_per_group: int, natoms: int) → numpy.ndarray