diff --git a/PhysicsAnalysis/AnalysisCommon/LeptonTaggers/python/LeptonTaggersConfig.py b/PhysicsAnalysis/AnalysisCommon/LeptonTaggers/python/LeptonTaggersConfig.py
index 61aec12c98f887ade3e4218b92906d248e348f09..f763377d8cbdf3ea7a0d589cd69bb93729408444 100644
--- a/PhysicsAnalysis/AnalysisCommon/LeptonTaggers/python/LeptonTaggersConfig.py
+++ b/PhysicsAnalysis/AnalysisCommon/LeptonTaggers/python/LeptonTaggersConfig.py
@@ -66,10 +66,23 @@ def GetDecoratePromptTauAlgs():
     return algs
 
 #------------------------------------------------------------------------------
-def GetExtraPromptVariablesForDxAOD(name='', addSpectators=False):
+def GetExtraPromptVariablesForDxAOD(name='', addSpectators=False, onlyBDT=True):
 
     prompt_lep_vars = []
 
+    #
+    # Decorate lepton only with the BDT outputs when the onlyBDT flag is true.
+    #
+    if onlyBDT:
+        if name == "" or name == "Electrons":
+            prompt_lep_vars += ["Electrons.PromptLeptonVeto.PromptLeptonIso."]
+
+        if name == "" or name == "Muons":
+            prompt_lep_vars += ["Muons.PromptLeptonVeto.PromptLeptonIso."]
+
+        return prompt_lep_vars
+ 
+ 
     prompt_vars  = "PromptLeptonVeto.PromptLeptonIso."
     prompt_vars += "PromptLeptonInput_TrackJetNTrack.PromptLeptonInput_sv1_jf_ntrkv."
     prompt_vars += "PromptLeptonInput_ip2.PromptLeptonInput_ip3."
@@ -104,10 +117,23 @@ def GetExtraPromptVariablesForDxAOD(name='', addSpectators=False):
 
 
 #------------------------------------------------------------------------------
-def GetExtraImprovedPromptVariablesForDxAOD(name=''):
+def GetExtraImprovedPromptVariablesForDxAOD(name='', onlyBDT=False):
 
     prompt_lep_vars = []
 
+    #
+    # Decorate lepton only with the BDT outputs when the onlyBDT flag is true.
+    #
+    if onlyBDT:
+        if name == "" or name == "Electrons":
+            prompt_lep_vars += ["Electrons.PromptLeptonImprovedVetoBARR.PromptLeptonImprovedVetoECAP."]
+
+        if name == "" or name == "Muons":
+            prompt_lep_vars += ["Muons.PromptLeptonImprovedVeto."]
+
+        return prompt_lep_vars
+ 
+
     prompt_vars  = "PromptLeptonImprovedInput_MVAXBin."
     prompt_vars += "PromptLeptonImprovedInput_PtFrac.PromptLeptonImprovedInput_DRlj."
     prompt_vars += "PromptLeptonImprovedInput_topoetcone30rel.PromptLeptonImprovedInput_ptvarcone30rel."
diff --git a/PhysicsAnalysis/DerivationFramework/DerivationFrameworkSM/share/STDM3.py b/PhysicsAnalysis/DerivationFramework/DerivationFrameworkSM/share/STDM3.py
index ebd85c416cd70c854f37d92f4f7877b02b40e324..77ea216df0b2e302f5f49b88adf9fa769671b9d9 100644
--- a/PhysicsAnalysis/DerivationFramework/DerivationFrameworkSM/share/STDM3.py
+++ b/PhysicsAnalysis/DerivationFramework/DerivationFrameworkSM/share/STDM3.py
@@ -213,8 +213,8 @@ reducedJetList = ["AntiKt2PV0TrackJets", "AntiKt4PV0TrackJets", "AntiKt4TruthJet
 replaceAODReducedJets(reducedJetList, STDM3Sequence, "STDM3Jets")
 
 # FAKE LEPTON TAGGER
-import JetTagNonPromptLepton.JetTagNonPromptLeptonConfig as JetTagConfig
-STDM3Sequence += JetTagConfig.GetDecoratePromptLeptonAlgs()
+import LeptonTaggers.LeptonTaggersConfig as LepTagConfig
+STDM3Sequence += LepTagConfig.GetDecorateImprovedPromptLeptonAlgs()
 
 # ADD SEQUENCE TO JOB
 DerivationFrameworkJob += STDM3Sequence
@@ -287,7 +287,7 @@ STDM3SlimmingHelper.IncludeMuonTriggerContent = True
 
 STDM3SlimmingHelper.ExtraVariables = ExtraContentAll
 STDM3SlimmingHelper.ExtraVariables += ["AntiKt4EMTopoJets.JetEMScaleMomentum_pt.JetEMScaleMomentum_eta.JetEMScaleMomentum_phi.JetEMScaleMomentum_m"]
-STDM3SlimmingHelper.ExtraVariables += JetTagConfig.GetExtraPromptVariablesForDxAOD()
+STDM3SlimmingHelper.ExtraVariables += LepTagConfig.GetExtraImprovedPromptVariablesForDxAOD(onlyBDT=True) 
 
 STDM3SlimmingHelper.AllVariables = ExtraContainersAll