diff --git a/mindspore/ccsrc/frontend/optimizer/ad/dfunctor.cc b/mindspore/ccsrc/frontend/optimizer/ad/dfunctor.cc
index e2147dc30e..bbff6fbec3 100644
--- a/mindspore/ccsrc/frontend/optimizer/ad/dfunctor.cc
+++ b/mindspore/ccsrc/frontend/optimizer/ad/dfunctor.cc
@@ -37,7 +37,7 @@ namespace ad {
 std::unordered_map<FuncGraphPtr, DFunctorPtr> DFunctor::func_graph_to_functor_;
 std::unordered_map<AnfNodePtr, AdjointPtr> DFunctor::anfnode_to_adjoin_definition_;
 
-bool lift_fv_before_grad = true;
+int lift_fv_before_grad = -1;
 
 DFunctor::DFunctor(const FuncGraphPtr &primal_graph, const pipeline::ResourceBasePtr &resources)
     : primal_graph_(primal_graph), resources_(resources), need_cut_(false), is_top_(false) {
@@ -75,8 +75,10 @@ void DFunctor::Clear() {
 }
 
 void DFunctor::BackPropagateFv(const AnfNodePtr &fv, const AnfNodePtr &din) {
-  if (lift_fv_before_grad) {
-    MS_LOG(EXCEPTION) << "BackPropagateFv can not find adjoint in anfnode_to_adjoin_ fv "
+  MS_EXCEPTION_IF_NULL(fv);
+  if (lift_fv_before_grad == 1) {
+    MS_EXCEPTION_IF_NULL(fv->func_graph());
+    MS_LOG(EXCEPTION) << "BackPropagateFv can not find adjoint in anfnode_to_adjoin_ fv:"
                       << fv->func_graph()->ToString() << " " << fv->ToString() << ".";
   }
   auto fv_adjoint = anfnode_to_adjoin_.find(fv);
@@ -444,7 +446,7 @@ AnfNodePtr DFunctor::AttachFvDoutToTape(const AnfNodePtr &grad_fv) {
   // Add grads wrt fv.
   const auto &free_variables_nodes = primal_graph_->free_variables_nodes();
   if (!is_top_ && free_variables_nodes.size() != 0) {
-    if (lift_fv_before_grad) {
+    if (lift_fv_before_grad == 1) {
       MS_LOG(EXCEPTION) << "direct fv size is: " << free_variables_nodes.size() << " in " << primal_graph_->ToString()
                         << ".";
     }
@@ -473,7 +475,7 @@ AnfNodePtr DFunctor::AttachFvDoutToTape(const AnfNodePtr &grad_fv) {
 }
 
 AnfNodePtr DFunctor::AttachIndirectFvDoutToTape(const AnfNodePtr &grad_fv) {
-  if (lift_fv_before_grad) {
+  if (lift_fv_before_grad == 1) {
     MS_LOG(EXCEPTION) << "Lift free variable case: AttachIndirectFvDoutToTape backprop indirect fv "
                       << grad_fv->ToString() << " " << primal_graph_->ToString() << ".";
   }
@@ -515,7 +517,7 @@ void DFunctor::MapMorphism() {
 
   // Set output for tape closure.
   AnfNodePtr grad_fv;
-  if (lift_fv_before_grad) {
+  if (lift_fv_before_grad == 1) {
     grad_fv = AttachFvDoutToTape(NewValueNode(newenv));
   } else {
     grad_fv = AttachIndirectFvDoutToTape(AttachFvDoutToTape(NewValueNode(newenv)));
diff --git a/mindspore/ccsrc/frontend/optimizer/ad/dfunctor.h b/mindspore/ccsrc/frontend/optimizer/ad/dfunctor.h
index 7aca3e21b0..1d0a111271 100644
--- a/mindspore/ccsrc/frontend/optimizer/ad/dfunctor.h
+++ b/mindspore/ccsrc/frontend/optimizer/ad/dfunctor.h
@@ -44,7 +44,7 @@ class DFunctor;
 using DFunctorPtr = std::shared_ptr<DFunctor>;
 
 // Flag to control if fv should be lifted before grad. If this lift_fv feature is mature, then this flag can be removed.
-extern bool lift_fv_before_grad;
+extern int lift_fv_before_grad;
 
 // D Functor's rules to map closure object and morphisms.
 class DFunctor : public std::enable_shared_from_this<DFunctor> {
diff --git a/mindspore/ccsrc/frontend/optimizer/ad/grad.cc b/mindspore/ccsrc/frontend/optimizer/ad/grad.cc
index 3eaeb2de5a..f6102bf016 100644
--- a/mindspore/ccsrc/frontend/optimizer/ad/grad.cc
+++ b/mindspore/ccsrc/frontend/optimizer/ad/grad.cc
@@ -60,6 +60,16 @@ FuncGraphPtr LiftFv(const pipeline::ResourceBasePtr &resource, const FuncGraphPt
   }
   return opt_fg;
 }
+
+bool NeedLiftFv(const FuncGraphPtr &func_graph) {
+  size_t switch_num = 0;
+  for (auto node : func_graph->manager()->all_nodes()) {
+    if (IsPrimitiveCNode(node, prim::kPrimSwitch)) {
+      ++switch_num;
+    }
+  }
+  return switch_num > 1;
+}
 }  // namespace
 
 FuncGraphPtr Grad(const FuncGraphPtr &func_graph, const pipeline::ResourceBasePtr &resources, bool is_top) {
@@ -74,8 +84,11 @@ FuncGraphPtr Grad(const FuncGraphPtr &func_graph, const pipeline::ResourceBasePt
   manager_ptr->AddFuncGraph(func_graph);
 
   FuncGraphPtr grad_fg = func_graph;
-  lift_fv_before_grad = (common::GetEnv("ENV_DONT_LIFT_FV_BEFORE_GRAD") != "1");
-  if (func_graph->func_graphs_used().size() != 0 && lift_fv_before_grad) {
+  // Only calculate lift_fv_before_grad once.
+  if (lift_fv_before_grad == -1) {
+    lift_fv_before_grad = (common::GetEnv("ENV_DONT_LIFT_FV_BEFORE_GRAD") != "1") && NeedLiftFv(func_graph) ? 1 : 0;
+  }
+  if (lift_fv_before_grad == 1 && func_graph->func_graphs_used().size() != 0) {
     grad_fg = LiftFv(resources, func_graph);
   }
   auto multi_graph_sink = [&func_graph](const FuncGraphPtr &f) {