torchref.refinement.optimizers.momentum_sa module

class torchref.refinement.optimizers.momentum_sa.MomentumStochasticSA(params, lr=0.001, betas=(0.9, 0.999), eps=1e-08, T_initial=1.0, T_final=0.01, total_steps=1000)[source]

Bases: Adam

Adam-based SA where noise is scaled by the adaptive learning rate, giving automatic scale invariance across parameters.

__init__(params, lr=0.001, betas=(0.9, 0.999), eps=1e-08, T_initial=1.0, T_final=0.01, total_steps=1000)[source]
step(closure=None)[source]

Perform a single optimization step.

Args:
closure (Callable, optional): A closure that reevaluates the model

and returns the loss.