@ptrblck: thanks for the link.
Yes, I was following that tutorial, and end up with the following that so far, work as expected.
I would be grateful if you could comment on it, thank you 
class BareNet(torch.nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super(BareNet, self).__init__()

        def _init_linear_layer(w, b):
            # Following:
            # https://github.com/pytorch/pytorch/blob/769cb5a6405b39a0678e6bc4f2d6fea62e0d3f12/torch/nn/modules/linear.py#L48
            stdv = 1. / math.sqrt(w.size(1))
            w.data.uniform_(-stdv, stdv)
            b.data.uniform_(-stdv, stdv)

        self.hidden_w = torch.nn.Parameter(torch.zeros(hidden_dim, input_dim))
        self.hidden_b = torch.nn.Parameter(torch.zeros(hidden_dim))
        self.hidden_pw = torch.nn.Parameter(torch.zeros(input_dim, hidden_dim))
        self.hidden_pb = torch.nn.Parameter(torch.zeros(hidden_dim))
        self.output_w = torch.nn.Parameter(torch.zeros(input_dim, hidden_dim))
        self.output_b = torch.nn.Parameter(torch.zeros(output_dim))
        self.output_pw = torch.nn.Parameter(torch.zeros(hidden_dim, input_dim))
        self.output_pb = torch.nn.Parameter(torch.zeros(output_dim))
        self.alpha = torch.nn.Parameter(torch.zeros(1))

        self.w_params = torch.nn.ParameterList([self.hidden_w, self.output_w])
        self.b_params = torch.nn.ParameterList([self.hidden_b, self.output_b])
        self.pw_params = torch.nn.ParameterList([self.hidden_pw, self.output_pw])
        self.pb_params = torch.nn.ParameterList([self.hidden_pb, self.output_pb])
        self.a_params = torch.nn.ParameterList([self.alpha])
        self.wb_params = [w for w in self.w_params] + [b for b in self.b_params]

        _init_linear_layer(self.hidden_w, self.hidden_b)
        _init_linear_layer(self.output_w, self.output_b)

    def forward(self, x):
        hidden_w = self.hidden_w.transpose(0, 1) + (self.alpha * self.hidden_pw)
        hidden_b = self.hidden_b + (self.alpha * self.hidden_pb)
        output_w = self.output_w.transpose(0, 1) + (self.alpha * self.output_pw)
        output_b = self.output_b + (self.alpha * self.output_pb)
        y = torch.nn.functional.relu( x.mm(hidden_w) + hidden_b )
        y = y.mm(output_w) + output_b
        return y

And in my_optim.py:
...

        def _phi(alpha):
            # Update a_, pw_, pb_ params
            for p in a_params: p.data.fill_(alpha)
            for i, p in enumerate(pw_params): p.data.copy_(w_step_dirs[i].transpose(0, 1))
            for i, p in enumerate(pb_params): p.data.copy_(b_step_dirs[i])

            # Get loss and its grad
            loss = closure(do_backward=False)
            grad_alpha, = torch.autograd.grad(loss, a_params, create_graph=False)

            # Zero a_, pw_, pb_ params
            for p in a_params: p.data.fill_(0.0)
            for p in pw_params: p.data.copy_(torch.zeros_like(p))
            for p in pb_params: p.data.copy_(torch.zeros_like(p))

            return (loss.item(), grad_alpha.item())
...