steps.lua

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

function loop(p, y, idxs)
  -- Only works if h is a derivable value as well
  x = p.x
  h = p.h
  for i = 1, x:size(1) do
    h[idxs[i]] = x[i]
  end
  return t.mean(t.pow(y - h, 2))
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])

dloop = grad(loop)  -- Crashes with {optimize = true}
d, loss = dloop({x=x_val, h=h_val}, y_val, idxs_val)
-- For some reason the gradient of x is empty
-- but h isn't, and that is simply the gradient in reverse
-- However, the last element seems to be 0, which is wrong?
print(d.h, loss)