Save self.numel() for backwards computation #5741
Description
Right now, the following derivatives.yaml declaration:
- name: mean(Tensor self)
self: grad.expand(self.sizes()) / self.numel()
produces the following backwards:
variable_list MeanBackward1::apply(const variable_list& grads) {
IndexRangeGenerator gen;
auto self_ix = gen.range(1);
variable_list grad_inputs(gen.size());
auto& grad = grads[0];
auto self = self_.unpack();
if (should_compute_output({ self_ix })) {
auto grad_result = grad.expand(self_sizes) / self.numel();
copy_range(grad_inputs, self_ix, grad_result);
}
return grad_inputs;
}
It looks like self.sizes() is being saved for the backward pass, but self.numel() isn't, requiring saving self for backwards.