steps.lua

t = require 'torch'
grad = require 'autograd'

function loop(x, h, y, idxs)
  for i = 1, x:size(1) do
    h[idxs[i]] = x[i]
    if i == 1 then
      cost = t.pow(y[idxs[i]] - h[idxs[i]], 2)
    else
      cost = cost + t.pow(y[idxs[i]] - h[idxs[i]], 2)
    end
  end
  return cost / x:size(1)
end

x_val = t.linspace(0, 1, arg[1])
y_val = t.linspace(0, 1, arg[1])
idxs_val = t.range(arg[1], 1, -1)
h_val = t.zeros(arg[1]):split(1)

dloop = grad(loop)
dx, loss = dloop(x_val, h_val, y_val, idxs_val)
print(dx, loss)