Assume the code looks like the following:
torch::autograd::variable_list SomeFunc::forward(torch::autograd::AutogradContext *ctx, torch::Tensor x) {
  /* ... */
  auto stream = c10::cuda::getCurrentCUDAStream();
  if(torch::GradMode::is_enabled()){
    // This will perform the forward computation in a different way using 
    // another slightly slower (in terms of inference only) formula having 
    // subexpressions that are in common with the backward computation, 
    // which makes the forward+backward fast.
    my_cuda_kernels::forward_and_precompute_for_backward(x, out, quantities_for_backward, stream);
    ctx->saved_data["quantities_for_backward"] = quantities_for_backward;
  } else {
    // The following will take a much faster path,
    // assumes we only need to do inference.
    my_cuda_kernels::forward_only(x, out, stream);
  }          
  return {out};
}

torch::autograd::variable_list SomeFunc::backward(torch::autograd::AutogradContext *ctx, torch::autograd::variable_list grad_output) {
  auto quantities_for_backward = ctx->saved_data["quantities_for_backward"].toTensor();
  /* do the backward computation by using "quantities_for_backward" to avoid recomputing some subexpressions that are common with the forward computation*/
}

The problem arises when I use this Function with autograd enabled (ie for example SomeFunc::apply(some_input).mean().backward() assuming some_input has requires_grad set to true), as during runtime libtorch will complain about getting a None instead of a Tensor when arriving at ctx->saved_data["quantities_for_backward"].toTensor();. So I assume that during the forward phase torch::GradMode::is_enabled() is not returning true.
When I remove the if statement in my forward method and just keep the version that precomputes the backward quantities and saves them in ctx in all cases, the code works again, which further confirms that torch::GradMode::is_enabled() is not working as intended inside a torch::autograd::Function.