def METHOD_NAME(self, train_x, train_y, xi, w_optim):
    """
    Compute unrolled weight w' (virtual step)
    Step process:
    1) forward
    2) calculate loss
    3) compute gradient (by backprop)
    4) update gradient
    Args:
        xi: learning rate for virtual gradient step (same as weights lr)
        w_optim: weights optimizer
    """
    # Forward and calculate loss
    # Loss for train with w. L_train(w)
    loss = self.model.loss(train_x, train_y)
    # Compute gradient
    gradients = torch.autograd.grad(loss, self.model.getWeights())
    # Do virtual step (Update gradient)
    # Below operations do not need gradient tracking
    with torch.no_grad():
        # dict key is not the value, but the pointer. So original network weight have to
        # be iterated also.
        for w, vw, g in zip(self.model.getWeights(), self.v_model.getWeights(), gradients):
            m = w_optim.state[w].get("momentum_buffer", 0.) * self.w_momentum
            vw.copy_(w - torch.FloatTensor(xi) * (m + g + self.w_weight_decay * w))
        # Sync alphas
        for a, va in zip(self.model.getAlphas(), self.v_model.getAlphas()):
            va.copy_(a)