Method __init__

tests/test_optimizer/test_adam_kernel.py:92–98

(self, lr: float, beta1: float, beta2: float, eps: float, weight_decay: float, use_adamw: bool)

class CPUAdamKernel(AdamKernel):
    def __init__(self, lr: float, beta1: float, beta2: float, eps: float, weight_decay: float, use_adamw: bool) -> None:
        super().__init__(lr, beta1, beta2, eps, weight_decay, use_adamw)
        from colossalai.kernel.kernel_loader import CPUAdamLoader

        cpu_optim = CPUAdamLoader().load()

        self.cpu_adam_op = cpu_optim.CPUAdamOptimizer(lr, beta1, beta2, eps, weight_decay, use_adamw)

    def update(self, step: int, param: Tensor, grad: Tensor, exp_avg: Tensor, exp_avg_sq: Tensor):
        self.cpu_adam_op.step(