def _optimizer_step(self, optimizer: Optimizer) -> None:
"""CPU and GPU optimization step.
Args:
optimizer: optimizer object
"""
optimizer.step()
def set_optimizer_momentum(optimizer: Optimizer, value: float, index: int = 0):
"""Set momentum of ``index`` 'th param group of optimizer to ``value``.
Args:
optimizer: PyTorch optimizer
value: new value of momentum
index (int, optional): integer index of optimizer's param groups,
default is 0
"""
betas = optimizer.param_groups[0].get("betas", None)
momentum = optimizer.param_groups[0].get("momentum", None)
if betas is not None:
_, beta = betas
optimizer.param_groups[index]["betas"] = (value, beta)
elif momentum is not None:
optimizer.param_groups[index]["momentum"] = value
def optimizer_step(self, loss: torch.Tensor, model: Model,
optimizer: Optimizer) -> None:
"""Abstraction over ``optimizer.step()`` step."""
optimizer.step()