fix: reset next_functions when backward completes

Author: Chamberlain0w0

infini_train/src/autograd/function.cc

@@ -57,6 +57,7 @@ std::vector<std::shared_ptr<Tensor>> Function::Apply(const std::vector<std::shar
     auto output_tensors = Forward(input_tensors);
     SetupContext(input_tensors, output_tensors);
 
+    next_functions_.clear();
     bool output_requires_grad = false;
     for (int idx = 0; idx < input_tensors.size(); ++idx) {
         const auto &input_tensor = input_tensors[idx];
@@ -107,5 +108,6 @@ void Function::BackwardPartial(const std::shared_ptr<Tensor> &grad_output, int g
 void Function::ResetState() {
     grad_outputs_reached_ = 0;
     for (auto &tensor : grad_outputs_) { tensor.reset(); }
+    for (auto &[next_function, idx] : next_functions_) { next_function.reset(); }
 }
 } // namespace infini_train::autograd