deepmd.pt.optimizer.LKF#Classes#LKFOptimizerFunctions#distribute_indices(→ list[tuple[int, int]])Module Contents#deepmd.pt.optimizer.LKF.distribute_indices(total_length: int, num_workers: int) → list[tuple[int, int]][source]#class deepmd.pt.optimizer.LKF.LKFOptimizer(params: Any, kalman_lambda: float = 0.98, kalman_nue: float = 0.9987, block_size: int = 5120)[source]#Bases: torch.optim.optimizer.Optimizer _params[source]# _state[source]# dist_init[source]# rank[source]# dindex = [][source]# remainder = 0[source]# __init_P() → None[source]# __get_blocksize() → int[source]# __get_nue() → float[source]# __split_weights(weight: torch.Tensor) → list[torch.Tensor][source]# __update(H: torch.Tensor, error: torch.Tensor, weights: torch.Tensor) → None[source]# set_grad_prefactor(grad_prefactor: float) → None[source]# step(error: torch.Tensor) → None[source]# get_device_id(index: int) → int | None[source]#
