Hey, thanks for your reply. Indeed, the m loop was in the wrong place. I’m going to try the implementation below. I only need to figure out how to do a single SGD momentum step when the number of iterations is still below the number of warmup steps.
Thanks!
    @torch.no_grad()
    def step(self, closure=None):

        for t in range(self.trajectory_param):

            self.zero_grad()

            loss = None
            if closure is not None:
                with torch.enable_grad():
                    loss = closure()

            for group in self.param_groups:

                # get parameters
                lr = group["lr"]
                alpha = group["alpha"]
                weight_decay = group["weight_decay"]
                cv_flag = group['cv_flag']

                for parameter in group["params"]:

                    if parameter.grad is None:
                        continue

                    state = self.state[parameter]
                    if len(state) == 0:
                        state["iteration"] = 0
                        state["momentum"] = torch.zeros_like(parameter)
                        if cv_flag:
                            parameter.total_grad = torch.zeros_like(parameter)
                            parameter.map_grad = torch.zeros_like(parameter)
                    state["iteration"] += 1

                    # get v_k 
                    momentum = state["momentum"]

                    # grad
                    gradient = parameter.grad
                    if weight_decay!=0.0:
                        gradient.add_(parameter, alpha=weight_decay)

                    if state["iteration"] > self.num_warmup_steps:

                        if self.resample_momentum:
                            state["momentum"] = gradient.new(gradient.size()).normal_(mean=0,std=1)
                            #state["momentum"] = torch.normal(mean=torch.zeros_like(gradient), std=1.0)

                        sigma = np.sqrt(alpha)
                        #w_scaled ~ N(0,2*alpha) / sqrt(lr), so that lr * w_scaled = N(0,alpha*lr)
                        sample_t = parameter.new(parameter.size()).normal_(mean=0.0, std=sigma) / np.sqrt(lr)
                        #v_k+1 = (1-alpha)*v_k - lr*grad + sqrt(alpha*lr)*w
                        if cv_flag:
                            gradient.add_(parameter.map_grad, alpha=-1.0)
                            gradient.add_(parameter.total_grad, alpha=1.0)
                        momentum_t = momentum.mul_(1.0-alpha).add_(0.5*gradient - sample_t, alpha=-lr)
                        #theta_k+1 = theta_k + v_k
                        parameter.add_(momentum_t)
                    else:
                        # single SGD momentum step if still in burnin stage
                        momentum_t = momentum.mul_(1.0-alpha).add_(0.5*gradient, alpha=-lr)
                        parameter.add_(momentum_t)
                        #break ?

        return loss