@cjekel
You don’'t need .to(out), torch::ones(1, 1) makes the grad output tensor size to 1 with initial value 1, that’s enough.  your original torch::ones_like(out) will end up with a tensor with ten 1s, so your calculated gradients will be sum up which will be 10 times bigger.