diff --git a/.gitignore b/.gitignore
index 80badc4c5fa860c5e3a166dea057c3c6d5fd8d7a..70cb36b8882ac80cbfb8d09a6fea5bebc74e87cf 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,4 +1,5 @@
 __pycache__
+.pytest_cache
 *.pyc
 *.pdf
 *.png
diff --git a/environment.yml b/environment.yml
index 1bcd59b6b6114612485c99a7fd077eea42c4e675..bb93a211fba54c20ec5015ae5be922bf51f36ef5 100644
--- a/environment.yml
+++ b/environment.yml
@@ -14,7 +14,9 @@ dependencies:
   - pip:
     - holidays=0.24
     - prophet=1.0.1
-    - pystan
+    - convertdate
+    - lunarcalendar
+    - pystan=2.19.1.1
     - hipe4ml
     - hipe4ml-converter
     - optuna
diff --git a/scripts/configs/DCM/test_config.json b/scripts/configs/DCM/test_config.json
index 0f578ac10ba087d80ff6c3546ed00c645c3426ef..2a5732226b3bc0e119565ff1adf3bd848e74845c 100644
--- a/scripts/configs/DCM/test_config.json
+++ b/scripts/configs/DCM/test_config.json
@@ -1 +1,85 @@
-{"file_names": {"training": "/lustre/cbm/users/tfic/pid_plain_trees/trees/PlainTree50K_DCM_trdrichrec_12agev.root", "test": "/lustre/cbm/users/tfic/pid_plain_trees/trees/PlainTree5K_DCM_trdrichrec_12agev.root"}, "var_names": {"momentum": "Complex_p", "charge": "Complex_q", "mass2": "Complex_mass2", "pid": "Complex_pid"}, "vars_to_draw": ["Complex_mass2", "Complex_p"], "features_for_train": ["Complex_pT", "Complex_mass2", "Complex_n_hits_trd", "Complex_e_loss_0", "Complex_e_loss_1", "Complex_e_loss_2", "Complex_e_loss_3", "Complex_vtx_chi2", "Complex_chi2_ov_ndf_vtx", "Complex_chi2_ov_ndf_trd", "Complex_M"], "cuts": {"Complex_mass2": {"lower": -1.0, "upper": 2.0}, "Complex_pT": {"lower": 0.0, "upper": 2.0}, "Complex_p": {"lower": -12.0, "upper": 12.0}, "Complex_eta": {"lower": 0.0, "upper": 6.0}}, "hyper_params": {"values": {"n_estimators": 776, "max_depth": 7, "learning_rate": 0.08112714035375232, "gamma": 9, "subsample": 0.71582, "alpha": 17}, "ranges": {"n_estimators": [300, 1200], "max_depth": [2, 8], "learning_rate": [0.01, 0.1], "gamma": [0, 10], "subsample": [0.3, 0.9], "alpha": [0, 20]}}, "bins": [0.38661597183346746, 1.4825726561740278, 2.0401938866642655, 2.7015859877268475, 3.7153128963006865, 12.0]}
\ No newline at end of file
+{
+    "file_names": {
+        "training": "/lustre/cbm/users/tfic/pid_plain_trees/trees/PlainTree50K_DCM_trdrichrec_12agev.root",
+        "test": "/lustre/cbm/users/tfic/pid_plain_trees/trees/PlainTree5K_DCM_trdrichrec_12agev.root"
+    },
+    "var_names": {
+        "momentum": "Complex_p",
+        "charge": "Complex_q",
+        "mass2": "Complex_mass2",
+        "pid": "Complex_pid"
+    },
+    "vars_to_draw": [
+        "Complex_mass2",
+        "Complex_p"
+    ],
+    "features_for_train": [
+        "Complex_pT",
+        "Complex_mass2",
+        "Complex_vtx_chi2",
+        "Complex_chi2_ov_ndf_vtx",
+        "Complex_M"
+    ],
+    "cuts": {
+        "Complex_mass2": {
+            "lower": -1.0,
+            "upper": 2.0
+        },
+        "Complex_pT": {
+            "lower": 0.0,
+            "upper": 2.0
+        },
+        "Complex_p": {
+            "lower": -12.0,
+            "upper": 12.0
+        },
+        "Complex_eta": {
+            "lower": 0.0,
+            "upper": 6.0
+        }
+    },
+    "hyper_params": {
+        "values": {
+            "n_estimators": 776,
+            "max_depth": 7,
+            "learning_rate": 0.08112714035375232,
+            "gamma": 9,
+            "subsample": 0.71582,
+            "alpha": 17
+        },
+        "ranges": {
+            "n_estimators": [
+                300,
+                1200
+            ],
+            "max_depth": [
+                2,
+                8
+            ],
+            "learning_rate": [
+                0.01,
+                0.1
+            ],
+            "gamma": [
+                0,
+                10
+            ],
+            "subsample": [
+                0.3,
+                0.9
+            ],
+            "alpha": [
+                0,
+                20
+            ]
+        }
+    },
+    "bins": [
+        0.38661597183346746,
+        1.4825726561740278,
+        2.0401938866642655,
+        2.7015859877268475,
+        3.7153128963006865,
+        12.0
+    ]
+}
\ No newline at end of file
diff --git a/scripts/jobs/train_no_plots_job.sh b/scripts/jobs/train_no_plots_job.sh
new file mode 100755
index 0000000000000000000000000000000000000000..625c54465eeaec2bc32a76ed9218fccfea2f6685
--- /dev/null
+++ b/scripts/jobs/train_no_plots_job.sh
@@ -0,0 +1,18 @@
+#!/bin/bash
+
+RESULT_DIR=$1
+NBINS=$2
+CONFIG=$3
+
+eval "$(conda shell.bash hook)"
+conda activate cbm23
+
+#needed for pritning graphs as slurm doesn't find it automatically
+export FONTCONFIG_FILE=$CONDA_PREFIX/etc/fonts/fonts.conf
+export FONTCONFIG_PATH=$CONDA_PREFIX/etc/fonts/
+
+cd $RESULT_DIR
+INDEX=${SLURM_ARRAY_TASK_ID}
+MLPID_SRC_DIR=/lustre/cbm/users/$USER/pid_ml/src
+python $MLPID_SRC_DIR/train_model.py -c=$CONFIG --autobins=$NBINS >&training_output_${INDEX}.txt
+
diff --git a/scripts/onnx_launch.sh b/scripts/onnx_launch.sh
new file mode 100755
index 0000000000000000000000000000000000000000..5f5be8e889bdd2dda9e68df0228bbba7d4cc8612
--- /dev/null
+++ b/scripts/onnx_launch.sh
@@ -0,0 +1,35 @@
+
+#!/bin/bash
+CONFIG=$PWD/configs/DCM/test_config.json
+NBINS=5
+
+timestamp=$(date +"%Y%m%d_%H%M%S")
+NEW_DIR_NAME="onnx_${timestamp}"
+RESULT_DIR=/lustre/cbm/users/$USER/pid/$NEW_DIR_NAME
+mkdir -p $RESULT_DIR
+echo "created directory $RESULT_DIR"
+LOG_DIR=$RESULT_DIR/log
+mkdir -p $LOG_DIR
+mkdir -p $LOG_DIR/out
+mkdir -p $LOG_DIR/error
+
+echo "logs can be found at $LOG_DIR"
+
+MLPID_SRC_DIR=/lustre/cbm/users/$USER/pid_ml/
+
+# sbatch --job-name="autobin"\
+#          --partition main\
+#          --mem=32000\
+#          --output=$LOG_DIR/out/%j.out.log \
+#          --error=$LOG_DIR/error/%j.err.log \
+#          --wait\
+#          -- $PWD/jobs/autobin_job.sh $RESULT_DIR $NBINS $CONFIG
+
+sbatch --job-name="train-all" \
+        -t 6:00:00 \
+         --partition main\
+         --mem=32000\
+        --output=$LOG_DIR/out/%j.out.log \
+        --error=$LOG_DIR/error/%j.err.log \
+        --array=1-$NBINS\
+        -- $PWD/jobs/train_no_plots_job.sh $RESULT_DIR $NBINS $CONFIG
\ No newline at end of file
diff --git a/src/onnx_config_generator.py b/src/onnx_config_generator.py
new file mode 100644
index 0000000000000000000000000000000000000000..18603856e89bcdfd46397a84b495c0399f7633c3
--- /dev/null
+++ b/src/onnx_config_generator.py
@@ -0,0 +1,86 @@
+import argparse
+import os
+import sys
+from typing import List
+
+from util.json_tools import load_json_file, save_json_file
+
+
+class OnnxConfigGenerator:
+    def __init__(self, path_to_config: str, output_dir: str, anti_particles: bool = False):
+        self.config: dict = load_json_file(path_to_config)
+        self.output_dir = output_dir
+        self.anti_particles = anti_particles
+
+        os.chdir(output_dir)
+        self.bins = self.config["bins"]
+        self.model_paths_list = []
+
+    def find_models(self):
+        for i in range(0, len(self.bins) - 1):
+            lo, hi = self.bins[i: i+2]
+            model_name = self.__construct_model_name(
+                lo, hi, self.anti_particles)
+            model_path = f'{self.output_dir}/{model_name}/{model_name}.onnx'
+            assert (os.path.exists(model_path))
+            model_dict = {
+                "lo": lo,
+                "hi": hi,
+                "path": model_path
+            }
+            self.model_paths_list.append(model_dict)
+
+    @staticmethod
+    def __construct_model_name(lo: float, hi: float, anti_particles: bool = False) -> str:
+        if anti_particles:
+            return f"model_{lo:.1f}_{hi:.1f}_anti"
+        else:
+            return f"model_{lo:.1f}_{hi:.1f}_positive"
+
+    def serialize(self):
+        output = {
+            "model_paths": self.model_paths_list,
+            "output_dir": self.output_dir,
+            "config_dir": self.config
+        }
+        save_json_file(output, f'{self.output_dir}/onnx_config.json')
+
+    @staticmethod
+    def parse_args(args: List[str]):
+        parser = argparse.ArgumentParser(
+            prog="Onnx config generator",
+            description="Making config for PID ML models so that they later load correctly during inferrence",
+        )
+        parser.add_argument(
+            "--config",
+            "-c",
+            required=True,
+            type=str,
+            help="Filename or path of config json file.",
+        )
+        parser.add_argument(
+            "--output_dir",
+            "-o",
+            required=True,
+            type=str,
+            help="Name of folder containing all folders with output models",
+        )
+        parser.add_argument(
+            "--antiparticles",
+            action="store_true",
+            help="If should train on particles instead of particles with positive charge.",
+        )
+        return parser.parse_args(args)
+
+
+if __name__ == "__main__":
+    args = OnnxConfigGenerator.parse_args(sys.argv[1:])
+
+    config_path = args.config
+    output_dir = args.output_dir
+
+    ocg = OnnxConfigGenerator(config_path, output_dir)
+    print("initialized")
+    ocg.find_models()
+    print("models found")
+    ocg.serialize()
diff --git a/src/train_model.py b/src/train_model.py
index 669dfc3b9ee0fd143c2a3bbee131267c7f54b578..cdd213c727b2660c751ef86e4e41d0c1ce12f77e 100755
--- a/src/train_model.py
+++ b/src/train_model.py
@@ -11,6 +11,7 @@ from typing import List
 import json
 
 from hipe4ml.model_handler import ModelHandler
+from hipe4ml_converter.h4ml_converter import H4MLConverter
 from sklearn.utils.class_weight import compute_sample_weight
 
 from util import json_tools, plotting_tools
@@ -23,9 +24,10 @@ class TrainModel:
     Class for training the ml model
     """
 
-    def __init__(self, model_hdl: ModelHandler, model_name: str):
+    def __init__(self, model_hdl: ModelHandler, model_name: str, config_file_name: str):
         self.model_hdl = model_hdl
         self.model_name = model_name
+        self.config_file_name = config_file_name
 
     def train_model_handler(
         self, train_test_data, sample_weights, model_hdl: ModelHandler = None
@@ -56,7 +58,13 @@ class TrainModel:
         model_name = model_name or self.model_name
         model_hdl = model_hdl or self.model_hdl
         model_hdl.dump_model_handler(model_name)
-        print(f"\nModel saved as {model_name}")
+        converter = H4MLConverter(model_hdl)
+        features_for_train = json_tools.load_features_for_train(
+            self.config_file_name)
+
+        converter.convert_model_onnx(len(features_for_train))
+        converter.dump_model_onnx(f"{model_name}.onnx")
+        print(f"\nModel saved as {model_name} and {model_name}.onnx")
 
 
 def parse_args(args: List[str]) -> argparse.Namespace:
@@ -223,7 +231,7 @@ if __name__ == "__main__":
         plotting_tools.opt_history_plot(study, save_plots)
         plotting_tools.opt_contour_plot(study, save_plots)
     # train model
-    train = TrainModel(model_hdl, model_name)
+    train = TrainModel(model_hdl, model_name, config_file_name)
     sample_weights = compute_sample_weight(
         # class_weight=None or {0: 1, 1: 3, 2: 1}, deleted for now
         class_weight="balanced",