}
    }
  }

 private:
  // TODO(b/174778737): Only hold needed inputs.
  vector<AbstractTensorHandle*> forward_inputs_;
};

class DivNoNanGradientFunction : public GradientFunction {
 public:
  explicit DivNoNanGradientFunction(vector<AbstractTensorHandle*> f_inputs,
                                    vector<AbstractTensorHandle*> f_outputs)
      : forward_inputs_(f_inputs), forward_outputs_(f_outputs) {