Commit e4642a66 authored by Manuel Guth's avatar Manuel Guth
Browse files

fixing umami train with nJets

parent b7016672
Pipeline #3226531 passed with stages
in 3 minutes and 20 seconds
......@@ -235,6 +235,12 @@ def Umami(args, train_config, preprocess_config):
logger.info(f"nJets: {nJets}, nTrks: {nTrks}")
logger.info(f"nFeatures: {nFeatures}, njet_features: {njet_features}")
if NN_structure["nJets_train"] is not None:
logger.info(
f"Training only on {NN_structure['nJets_train']} jets as specified in the config."
)
nJets = NN_structure["nJets_train"]
umami, _ = Umami_model(
train_config=train_config,
input_shape=(nTrks, nFeatures),
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment