} else {
        in_grads.push_back(nullptr);
      }
    } else {
      in_grads.push_back(current_grad->second);
    }
  }

  // Avoid infinite recursion. Whichever forward function we run, it'll end up
  // executing ops, and we don't want to watch those with this accumulator.
  call_state_.emplace(nullptr, true);
  auto pop_call_state = gtl::MakeCleanup([this] { this->call_state_.pop(); });