Method __init__

tests/test_optimizer/test_adam_kernel.py:66–72

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

class FusedAdamKernel(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 FusedOptimizerLoader

        fused_optim = FusedOptimizerLoader().load()
        self.fused_adam = fused_optim.multi_tensor_adam
        self.dummy_overflow_buf = torch.tensor([0], dtype=torch.int, device=get_accelerator().get_current_device())

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