N = 64;
D_in = 1000;
H = 100;
D_out = 10;

% init input and output
x = rand(N, D_in)-0.5;
y = rand(N, D_out)-0.5;

% init weights
w1 = rand(D_in, H)-0.5;
w2 = rand(H, D_out)-0.5;

lr = 1e-4;
total_loss = {};
for i = 1:500
% forward
h = x * w1;
h_relu = max(h, 0);
y_pred = h_relu * w2;
loss = norm(y_pred-y, 2);
disp(loss);
total_loss{i} = loss;
% backward
grad_y_pred = 2 * (y_pred-y);
grad_w2 = h_relu' * grad_y_pred;
grad_h_relu = grad_y_pred * w2';
grad_h = grad_h_relu .* (h>=0);
grad_w1 = x' * grad_h;
% optimizer
w1 = w1 - lr * grad_w1;
w2 = w2 - lr * grad_w2;
end

total_loss = cell2mat(total_loss)
plot(total_loss)