int64_t tensor_id = source_tensor_ids[i];
    auto grad_it = gradients.find(tensor_id);
    if (grad_it == gradients.end()) {
      result[i] = nullptr;
    } else {
      if (grad_it->second.size() > 1) {
        Gradient* grad = vspace.AggregateGradients(grad_it->second);
        grad_it->second.clear();
        grad_it->second.push_back(grad);
      }
      result[i] = grad_it->second[0];
      used_gradient_ids.insert(tensor_id);