diff --git a/docs/advanced/model_optimization.rst b/docs/advanced/model_optimization.rst
index 07afcfa48f..3c3caba7d5 100644
--- a/docs/advanced/model_optimization.rst
+++ b/docs/advanced/model_optimization.rst
@@ -109,9 +109,11 @@ Finally, optimizing Vivado DSPs is possible, given a hls4ml config:
     # Note the change from ParameterEstimator to VivadoDSPEstimator
     optimized_model = optimize_keras_for_hls4ml(
         baseline_model, model_attributes, VivadoDSPEstimator, scheduler,
-        X_train, y_train, X_val, y_val, batch_size, epochs, optimizer, loss_fn, metric, increasing, rtol
+        X_train, y_train, X_val, y_val, batch_size, epochs,
+        optimizer, loss_fn, metric, increasing, rtol
     )
 
+
 There are two more Vivado "optimizers" - VivadoFFEstimator, aimed at reducing register utilisation and VivadoMultiObjectiveEstimator, aimed at optimising BRAM and DSP utilisation.
 Note, to ensure DSPs are optimized, "unrolled" Dense multiplication must be used before synthesing HLS, by modifying the config:
 .. code-block:: Python
diff --git a/setup.cfg b/setup.cfg
index 8e0fdd4ba1..2f2fe1847d 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -41,9 +41,9 @@ pytest_randomly.random_seeder =
 
 [options.extras_require]
 optimization =
-    keras-surgeon@https://github.com/fastmachinelearning/keras-surgeon
-    ortools
-    packaging
+    keras-surgeon@git+https://github.com/fastmachinelearning/keras-surgeon.git
+    ortools==9.4.1874
+    profiling
 profiling =
     matplotlib
     pandas
diff --git a/test/pytest/generate_ci_yaml.py b/test/pytest/generate_ci_yaml.py
index 6d816a7abb..46c1fa9e5b 100644
--- a/test/pytest/generate_ci_yaml.py
+++ b/test/pytest/generate_ci_yaml.py
@@ -32,5 +32,26 @@ def uses_example_model(test_filename):
     else:
         yml.update(new_yml)
 
+# hls4ml Optimization API
+tests = glob.glob('test_optimization/test_*.py')
+for test in tests:
+    name = test.replace('test_optimization/', '').replace('test_', '').replace('.py', '')
+    new_yml = yaml.safe_load(template.format(name, f'test_optimization/test_{name}.py', int(uses_example_model(test))))
+    if yml is None:
+        yml = new_yml
+    else:
+        yml.update(new_yml)
+
+tests = glob.glob('test_optimization/test_keras/test_*.py')
+for test in tests:
+    name = test.replace('test_optimization/test_keras/', '').replace('test_', '').replace('.py', '')
+    new_yml = yaml.safe_load(
+        template.format(name, f'test_optimization/test_keras/test_{name}.py', int(uses_example_model(test)))
+    )
+    if yml is None:
+        yml = new_yml
+    else:
+        yml.update(new_yml)
+
 yamlfile = open('pytests.yml', 'w')
 yaml.safe_dump(yml, yamlfile)