diff --git a/Project.toml b/Project.toml
index 4e59372..e944f5c 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "LearningToOptimize"
 uuid = "e1d8bfa7-c465-446a-84b9-451470f6e76c"
 authors = ["andrewrosemberg <andrewrosemberg@gmail.com> and contributors"]
-version = "1.0.0"
+version = "1.1.0"
 
 [deps]
 Arrow = "69666777-d1a9-59fb-9406-91d4454c9d45"
@@ -11,6 +11,7 @@ Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c"
 JuMP = "4076af6c-e467-56ae-b986-b466b2749572"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+MLJ = "add582a8-e3ab-11e8-2d5e-e98b27df1bc7"
 MLJFlux = "094fc8d1-fd35-5302-93ea-dabda2abf845"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
@@ -25,14 +26,15 @@ Arrow = "2"
 CSV = "0.10"
 DataFrames = "1"
 Distributions = "0.25"
-Flux = "0.14"
+Flux = "0.14, 0.16"
 JuMP = "1"
 MLJFlux = "0.6"
+MLJ = "0.20"
 NNlib = "0.9"
-Optimisers = "0.3"
-ParametricOptInterface = "0.8"
+Optimisers = "0.3, 0.4"
+ParametricOptInterface = "0.8, 0.9"
 Statistics = "1"
-Zygote = "0.6.68"
+Zygote = "0.6.68, 0.7"
 julia = "1.9"
 
 [extras]
@@ -41,10 +43,9 @@ Clarabel = "61c947e1-3e6d-4ee4-985a-eec8c727bd6e"
 DelimitedFiles = "8bb1440f-4735-579b-a4ab-409b98df4dab"
 HiGHS = "87dc4568-4c63-4d18-b0c0-bb2238e4078b"
 Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9"
-MLJ = "add582a8-e3ab-11e8-2d5e-e98b27df1bc7"
 PGLib = "07a8691f-3d11-4330-951b-3c50f98338be"
 PowerModels = "c36e90e8-916a-50a6-bd94-075b64ef4655"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "DelimitedFiles", "PGLib", "HiGHS", "PowerModels", "Clarabel", "Ipopt", "MLJ"]
+test = ["Test", "DelimitedFiles", "PGLib", "HiGHS", "PowerModels", "Clarabel", "Ipopt"]
diff --git a/README.md b/README.md
index d3287c6..198242a 100644
--- a/README.md
+++ b/README.md
@@ -6,7 +6,7 @@
     </div>
 </div>
 
-Learning to optimize (LearningToOptimize) package that provides basic functionalities to help fit proxy models for optimization.
+Learning to optimize (LearningToOptimize) package that provides basic functionalities to help fit proxy models for parametric optimization problems.
 
 Have a look at our sister [HugginFace Organization](https://huggingface.co/LearningToOptimize), for datasets, pre-trained models and benchmarks.
 
@@ -19,6 +19,34 @@ Have a look at our sister [HugginFace Organization](https://huggingface.co/Learn
 
 ![flowchart](docs/src/assets/L2O.png)
 
+# Background
+
+Parametric optimization problems arise in scenarios where certain elements (e.g., coefficients, constraints) may vary according to problem parameters. A general form of a parameterized convex optimization problem is 
+
+$$
+\begin{aligned}
+&\min_{x} \quad f(x; \theta) \\
+&\text{subject to} \quad g_i(x; \theta) \leq 0, \quad i = 1,\dots, m \\
+&\quad\quad\quad\quad A(\theta)x = b(\theta)
+\end{aligned}
+$$
+
+where $ \theta $ is the parameter.
+
+**Learning to Optimize (L2O)** is an emerging paradigm where machine learning models *learn* to solve optimization problems efficiently. This approach is also known as using **optimization proxies** or **amortized optimization**. 
+
+In more technical terms, **amortized optimization** seeks to learn a function \\( f_\theta(x) \\) that maps problem parameters \\( x \\) to solutions \\( y \\) that (approximately) minimize a given objective function subject to constraints. Modern methods leverage techniques like **differentiable optimization layers**, **input-convex neural networks**, or constraint-enforcing architectures (e.g., [DC3](https://openreview.net/pdf?id=0Ow8_1kM5Z)) to ensure that the learned proxy solutions are both feasible and performant. By coupling the solver and the model in an **end-to-end** pipeline, these approaches let the training objective directly reflect downstream metrics, improving speed and reliability.
+
+Recent advances also focus on **trustworthy** or **certifiable** proxies, where constraint satisfaction or performance bounds are guaranteed. This is crucial in domains like energy systems or manufacturing, where infeasible solutions can have large penalties or safety concerns. Overall, learning-based optimization frameworks aim to combine the advantages of ML (data-driven generalization) with the rigor of mathematical programming (constraint handling and optimality).
+
+For a broader overview, see the [SIAM News article on trustworthy optimization proxies](https://www.siam.org/publications/siam-news/articles/fusing-artificial-intelligence-and-optimization-with-trustworthy-optimization-proxies/), which highlights the growing synergy between AI and classical optimization.
+
+# Installation
+
+```julia
+] add LearningToOptimize
+```
+
 ## Generate Dataset
 This package provides a basic way of generating a dataset of the solutions of an optimization problem by varying the values of the parameters in the problem and recording it.
 
@@ -62,7 +90,33 @@ Which creates the following CSV:
 |  9 | 9.0 |
 | 10 | 10.0|
 
-ps.: For illustration purpose, I have represented the id's here as integers, but in reality they are generated as UUIDs. 
+ps.: For illustration purpose, I have represented the id's here as integers, but in reality they are generated as UUIDs.
+
+To load the parameter values back:
+
+```julia
+problem_iterator = load("input_file.csv", CSVFile)
+```
+
+### Samplers
+
+Instead of defining parameter instances manually, one may sample parameter values using pre-defined samplers - e.g. `scaled_distribution_sampler`, `box_sampler`- or define their own sampler. Samplers are functions that take a vector of parameter of type `MOI.Parameter` and return a matrix of parameter values.
+
+The easiest way to go from problem definition, sampling parameter values and saving them is to use the `general_sampler` function: 
+
+```julia
+general_sampler(
+    "examples/powermodels/data/6468_rte/6468_rte_SOCWRConicPowerModel_POI_load.mof.json";
+    samplers = [
+        (original_parameters) -> scaled_distribution_sampler(original_parameters, 10000),
+        (original_parameters) -> line_sampler(original_parameters, 1.01:0.01:1.25),
+        (original_parameters) -> box_sampler(original_parameters, 300),
+    ],
+)
+```
+
+This function is a general sampler that uses a set of samplers to sample the parameter space. 
+It loads the underlying model from a passed `file` that works with JuMP's `read_from_file` (ps.: currently only tested with `MathOptFormat`), samples the parameters and saves the sampled parameters to `save_file`.
 
 ### The Recorder
 
@@ -104,13 +158,15 @@ recorder = Recorder{ArrowFile}("output_file.arrow", primal_variables=[x], dual_v
 In order to train models to be able to forecast optimization solutions from parameter values, one option is to use the package Flux.jl:
 
 ```julia
+using CSV, DataFrames, Flux
+
 # read input and output data
 input_data = CSV.read("input_file.csv", DataFrame)
 output_data = CSV.read("output_file.csv", DataFrame)
 
 # Separate input and output variables
-output_variables = output_data[!, Not(:id)]
-input_features = innerjoin(input_data, output_data[!, [:id]], on = :id)[!, Not(:id)] # just use success solves
+output_variables = output_data[!, Not([:id, :status, :primal_status, :dual_status, :objective, :time])] # just predict solutions
+input_features = innerjoin(input_data, output_data[!, [:id]]; on=:id)[!, Not(:id)] # just use success solves
 
 # Define model
 model = Chain(
@@ -136,6 +192,38 @@ Flux.train!(loss, Flux.params(model), [(input_features, output_variables)], opti
 predictions = model(input_features)
 ```
 
+Another option is to use the package MLJ.jl:
+
+```julia
+using MLJ
+
+# Define the model
+model = MultitargetNeuralNetworkRegressor(;
+    builder=FullyConnectedBuilder([64, 32]),
+    rng=123,
+    epochs=20,
+    optimiser=Optimisers.Adam(),
+)
+
+# Train the model
+mach = machine(model, input_features, output_variables)
+fit!(mach; verbosity=2)
+
+# Make predictions
+predict(mach, input_features)
+
+```
+
+### Evaluating the ML model
+
+For ease of use, we built a general evaluator that can be used to evaluate the model.
+It will return a `NamedTuple` with the objective value and infeasibility of the 
+predicted solution for each instance, and the overall inference time and allocated memory.
+
+```julia
+evaluation = general_evaluator(problem_iterator, mach)
+```
+
 ## Coming Soon
 
 Future features:
diff --git a/docs/make.jl b/docs/make.jl
index da5b4cf..ab398ee 100644
--- a/docs/make.jl
+++ b/docs/make.jl
@@ -1,24 +1,30 @@
 using LearningToOptimize
 using Documenter
 
-DocMeta.setdocmeta!(LearningToOptimize, :DocTestSetup, :(using LearningToOptimize); recursive=true)
+DocMeta.setdocmeta!(
+    LearningToOptimize,
+    :DocTestSetup,
+    :(using LearningToOptimize);
+    recursive = true,
+)
 
 makedocs(;
-    modules=[LearningToOptimize],
-    authors="andrewrosemberg <andrewrosemberg@gmail.com> and contributors",
-    repo="https://github.com/andrewrosemberg/LearningToOptimize.jl/blob/{commit}{path}#{line}",
-    sitename="LearningToOptimize.jl",
-    format=Documenter.HTML(;
-        prettyurls=get(ENV, "CI", "false") == "true",
-        canonical="https://andrewrosemberg.github.io/LearningToOptimize.jl",
-        edit_link="main",
-        assets=String[],
+    modules = [LearningToOptimize],
+    authors = "andrewrosemberg <andrewrosemberg@gmail.com> and contributors",
+    repo = "https://github.com/andrewrosemberg/LearningToOptimize.jl/blob/{commit}{path}#{line}",
+    sitename = "LearningToOptimize.jl",
+    format = Documenter.HTML(;
+        prettyurls = get(ENV, "CI", "false") == "true",
+        canonical = "https://andrewrosemberg.github.io/LearningToOptimize.jl",
+        edit_link = "main",
+        assets = String[],
     ),
-    pages=["Home" => "index.md",
+    pages = [
+        "Home" => "index.md",
         "Arrow" => "arrow.md",
         "Parameter Type" => "parametertype.md",
         "API" => "api.md",
     ],
 )
 
-deploydocs(; repo="github.com/andrewrosemberg/LearningToOptimize.jl", devbranch="main")
+deploydocs(; repo = "github.com/andrewrosemberg/LearningToOptimize.jl", devbranch = "main")
diff --git a/docs/src/index.md b/docs/src/index.md
index 47c5070..a259ed0 100644
--- a/docs/src/index.md
+++ b/docs/src/index.md
@@ -71,7 +71,33 @@ Which creates the following CSV:
 |  9 | 9.0 |
 | 10 | 10.0|
 
-ps.: For illustration purpose, I have represented the id's here as integers, but in reality they are generated as UUIDs. 
+ps.: For illustration purpose, I have represented the id's here as integers, but in reality they are generated as UUIDs.
+
+To load the parameter values back:
+
+```julia
+problem_iterator = load("input_file.csv", CSVFile)
+```
+
+### Samplers
+
+Instead of defining parameter instances manually, one may sample parameter values using pre-defined samplers - e.g. `scaled_distribution_sampler`, `box_sampler`- or define their own sampler. Samplers are functions that take a vector of parameter of type `MOI.Parameter` and return a matrix of parameter values.
+
+The easiest way to go from problem definition, sampling parameter values and saving them is to use the `general_sampler` function: 
+
+```julia
+general_sampler(
+    "examples/powermodels/data/6468_rte/6468_rte_SOCWRConicPowerModel_POI_load.mof.json";
+    samplers = [
+        (original_parameters) -> scaled_distribution_sampler(original_parameters, 10000),
+        (original_parameters) -> line_sampler(original_parameters, 1.01:0.01:1.25),
+        (original_parameters) -> box_sampler(original_parameters, 300),
+    ],
+)
+```
+
+This function is a general sampler that uses a set of samplers to sample the parameter space. 
+It loads the underlying model from a passed `file` that works with JuMP's `read_from_file` (ps.: currently only tested with `MathOptFormat`), samples the parameters and saves the sampled parameters to `save_file`.
 
 ### The Recorder
 
@@ -113,13 +139,15 @@ recorder = Recorder{ArrowFile}("output_file.arrow", primal_variables=[x], dual_v
 In order to train models to be able to forecast optimization solutions from parameter values, one option is to use the package Flux.jl:
 
 ```julia
+using CSV, DataFrames, Flux
+
 # read input and output data
 input_data = CSV.read("input_file.csv", DataFrame)
 output_data = CSV.read("output_file.csv", DataFrame)
 
 # Separate input and output variables
-output_variables = output_data[!, Not(:id)]
-input_features = innerjoin(input_data, output_data[!, [:id]], on = :id)[!, Not(:id)] # just use success solves
+output_variables = output_data[!, Not([:id, :status, :primal_status, :dual_status, :objective, :time])] # just predict solutions
+input_features = innerjoin(input_data, output_data[!, [:id]]; on=:id)[!, Not(:id)] # just use success solves
 
 # Define model
 model = Chain(
@@ -145,18 +173,40 @@ Flux.train!(loss, Flux.params(model), [(input_features, output_variables)], opti
 predictions = model(input_features)
 ```
 
-## Coming Soon
+Another option is to use the package MLJ.jl:
 
-Future features:
- - ML objectives that penalize infeasible predictions;
- - Warm-start from predicted solutions.
+```julia
+using MLJ
+
+# Define the model
+model = MultitargetNeuralNetworkRegressor(;
+    builder=FullyConnectedBuilder([64, 32]),
+    rng=123,
+    epochs=20,
+    optimiser=Optimisers.Adam(),
+)
+
+# Train the model
+mach = machine(model, input_features, output_variables)
+fit!(mach; verbosity=2)
 
+# Make predictions
+predict(mach, input_features)
 
-<!-- ```@index
+```
 
-``` -->
+### Evaluating the ML model
 
+For ease of use, we built a general evaluator that can be used to evaluate the model.
+It will return a `NamedTuple` with the objective value and infeasibility of the 
+predicted solution for each instance, and the overall inference time and allocated memory.
 
-<!-- ```@autodocs
-Modules = [LearningToOptimize]
-``` -->
+```julia
+evaluation = general_evaluator(problem_iterator, mach)
+```
+
+## Coming Soon
+
+Future features:
+ - ML objectives that penalize infeasible predictions;
+ - Warm-start from predicted solutions.
diff --git a/examples/cuttingplanes/cutting_planes.jl b/examples/cuttingplanes/cutting_planes.jl
index 0510d3b..b18bf0e 100644
--- a/examples/cuttingplanes/cutting_planes.jl
+++ b/examples/cuttingplanes/cutting_planes.jl
@@ -53,7 +53,7 @@ function copy_binary_constraint(::JuMP.Model, ::VariableRef, set, var_mapping)
     return nothing
 end
 function copy_binary_constraint(m_to::JuMP.Model, func::AffExpr, set, var_mapping)
-    terms_dict = JuMP.OrderedCollections.OrderedDict{VariableRef, Float64}()
+    terms_dict = JuMP.OrderedCollections.OrderedDict{VariableRef,Float64}()
     for var in keys(func.terms)
         terms_dict[var_mapping[var]] = func.terms[var]
     end
@@ -74,7 +74,7 @@ end
 # Copy objective function
 function copy_binary_objective(m_to::JuMP.Model, m_from::JuMP.Model, var_mapping::Dict)
     obj = objective_function(m_from)
-    obj2_dict = JuMP.OrderedCollections.OrderedDict{VariableRef, Float64}()
+    obj2_dict = JuMP.OrderedCollections.OrderedDict{VariableRef,Float64}()
     for var in keys(obj.terms)
         if is_binary(var)
             obj2_dict[var_mapping[var]] = obj.terms[var]
@@ -94,7 +94,7 @@ function copy_binary_model(m_from::JuMP.Model)
 end
 
 # remove binary terms
-function delete_binary_terms!(m::JuMP.Model; delete_objective=true)
+function delete_binary_terms!(m::JuMP.Model; delete_objective = true)
     if delete_objective
         obj = objective_function(m)
         for var in keys(obj.terms)
@@ -111,18 +111,19 @@ function delete_binary_terms!(m::JuMP.Model; delete_objective=true)
     for con in all_binary_constraints(m)
         delete(m, con)
     end
-    
+
     return
 end
 
-function add_deficit_constraints!(model::JuMP.Model; penalty=nothing)
+function add_deficit_constraints!(model::JuMP.Model; penalty = nothing)
     if isnothing(penalty)
         obj = objective_function(model)
         # get the highest coefficient
         penalty = maximum(abs.(values(obj.terms)))
         penalty = penalty * 1.1
     end
-    consrefs = [con for con in all_constraints(model, include_variable_in_set_constraints=false)]
+    consrefs =
+        [con for con in all_constraints(model, include_variable_in_set_constraints = false)]
     @variable(model, deficit[1:length(consrefs)])
     @variable(model, norm_deficit)
     for (i, eq) in enumerate(consrefs)
@@ -156,16 +157,17 @@ end
 
 # add cut
 function add_cut!(upper_model::JuMP.Model, cut_intercept, cut_slope, cut_point, u)
-    @constraint(upper_model, 
+    @constraint(
+        upper_model,
         upper_model[:θ] >= cut_intercept + dot(cut_slope, u .- cut_point)
     )
 end
 
 function solve_or_shuffle!(model::JuMP.Model, u, cuts_point, iteration)
     JuMP.optimize!(model)
-    
+
     if termination_status(model) != MOI.OPTIMAL
-        return rand([0;1],length(u))
+        return rand([0; 1], length(u))
     end
 
     new_point = round.(Int, value.(u))
@@ -173,14 +175,21 @@ function solve_or_shuffle!(model::JuMP.Model, u, cuts_point, iteration)
         return new_point
     end
 
-    if all(isapprox(new_point, cuts_point[j]) for j in iteration-10:iteration-1)
-        return rand([0;1],length(u))
+    if all(isapprox(new_point, cuts_point[j]) for j = iteration-10:iteration-1)
+        return rand([0; 1], length(u))
     end
 
     return new_point
 end
 
-function cutting_planes!(inner_model::JuMP.Model; upper_solver, inner_solver, max_iter::Int=4000, atol::Real=0.1, bound::Real=0.0)
+function cutting_planes!(
+    inner_model::JuMP.Model;
+    upper_solver,
+    inner_solver,
+    max_iter::Int = 4000,
+    atol::Real = 0.1,
+    bound::Real = 0.0,
+)
     upper_model, inner_2_upper_map, cons_mapping = copy_binary_model(inner_model)
     delete_binary_terms!(inner_model)
     add_deficit_constraints!(inner_model)
@@ -194,7 +203,7 @@ function cutting_planes!(inner_model::JuMP.Model; upper_solver, inner_solver, ma
     JuMP.optimize!(inner_model)
     cuts_intercept = [objective_value(inner_model)]
     cuts_slope = [[MOI.get(inner_model, POI.ParameterDual(), u_i) for u_i in u_inner]]
-    cuts_point = [rand([0;1],length(upper_2_inner))]
+    cuts_point = [rand([0; 1], length(upper_2_inner))]
 
     # cutting planes epigraph variable
     @variable(upper_model, θ >= bound)
@@ -209,27 +218,31 @@ function cutting_planes!(inner_model::JuMP.Model; upper_solver, inner_solver, ma
     while i <= max_iter
         # Add cuts
         add_cut!(upper_model, cuts_intercept[i], cuts_slope[i], cuts_point[i], u)
-    
+
         # Add point to the lists
         new_point = solve_or_shuffle!(upper_model, u, cuts_point, i)
         push!(cuts_point, new_point)
         if value(θ) > bound
             bound = value(θ)
         end
-    
+
         # run inner problem
         set_binary_variables!(inner_model, u_inner, cuts_point[i+1])
         JuMP.optimize!(inner_model)
-    
+
         # Add cut to the lists
-        if termination_status(inner_model) == MOI.OPTIMAL || termination_status(inner_model) == MOI.LOCALLY_SOLVED
+        if termination_status(inner_model) == MOI.OPTIMAL ||
+           termination_status(inner_model) == MOI.LOCALLY_SOLVED
             push!(cuts_intercept, objective_value(inner_model))
-            push!(cuts_slope, [MOI.get(inner_model, POI.ParameterDual(), u_i) for u_i in u_inner])
+            push!(
+                cuts_slope,
+                [MOI.get(inner_model, POI.ParameterDual(), u_i) for u_i in u_inner],
+            )
         else
             println("Inner problem failed")
-            break;
+            break
         end
-    
+
         # test convergence
         u_bound = cuts_intercept[i+1]
         upper_bound[i] = u_bound
@@ -237,7 +250,7 @@ function cutting_planes!(inner_model::JuMP.Model; upper_solver, inner_solver, ma
         gap[i] = abs(bound - u_bound) / u_bound
         if i > 10 && gap[i] < atol # && all([all(cuts_point[i] .== cuts_point[j]) for j in i-10:i-1])
             println("Converged")
-            break;
+            break
         else
             @info "Iteration $i" bound cuts_intercept[i]
         end
@@ -250,4 +263,3 @@ function cutting_planes!(inner_model::JuMP.Model; upper_solver, inner_solver, ma
     lower_bound = lower_bound[1:i]
     return upper_model, lower_bound, upper_bound, gap
 end
-
diff --git a/examples/generate_dataset_inputs.jl b/examples/generate_dataset_inputs.jl
index 214f57d..c6e5502 100644
--- a/examples/generate_dataset_inputs.jl
+++ b/examples/generate_dataset_inputs.jl
@@ -2,9 +2,9 @@ using LearningToOptimize
 
 general_sampler(
     "examples/powermodels/data/6468_rte/6468_rte_SOCWRConicPowerModel_POI_load.mof.json";
-    samplers=[
+    samplers = [
         (original_parameters) -> scaled_distribution_sampler(original_parameters, 10000),
-        (original_parameters) -> line_sampler(original_parameters, 1.01:0.01:1.25), 
+        (original_parameters) -> line_sampler(original_parameters, 1.01:0.01:1.25),
         (original_parameters) -> box_sampler(original_parameters, 300),
     ],
-)
\ No newline at end of file
+)
diff --git a/examples/powermodels/HuggingFaceDatasets.jl b/examples/powermodels/HuggingFaceDatasets.jl
index d6a4f92..2b6eaf9 100644
--- a/examples/powermodels/HuggingFaceDatasets.jl
+++ b/examples/powermodels/HuggingFaceDatasets.jl
@@ -4,9 +4,16 @@ using Conda
 
 huggingface_hub = pyimport("huggingface_hub")
 
-huggingface_hub.login(token=ENV["HUGGINGFACE_TOKEN"])
+huggingface_hub.login(token = ENV["HUGGINGFACE_TOKEN"])
 
-function download_dataset(organization, dataset, case_name, io_type; formulation="", cache_dir="./data/")
+function download_dataset(
+    organization,
+    dataset,
+    case_name,
+    io_type;
+    formulation = "",
+    cache_dir = "./data/",
+)
     dataset_url = joinpath(organization, dataset)
     if io_type ∉ ["input", "output"]
         throw(ArgumentError("io_type must be 'input' or 'output'."))
@@ -22,18 +29,31 @@ function download_dataset(organization, dataset, case_name, io_type; formulation
     end
 
     # Fetch the dataset from the provided URL
-    huggingface_hub.snapshot_download(dataset_url, allow_patterns=["$data_path/*.arrow"], local_dir=cache_dir, repo_type="dataset", local_dir_use_symlinks=false)
-    
+    huggingface_hub.snapshot_download(
+        dataset_url,
+        allow_patterns = ["$data_path/*.arrow"],
+        local_dir = cache_dir,
+        repo_type = "dataset",
+        local_dir_use_symlinks = false,
+    )
+
     return nothing
 end
 
-cache_dir="./examples/powermodels/data/"
+cache_dir = "./examples/powermodels/data/"
 organization = "LearningToOptimize"
 dataset = "pglib_opf_solves"
 case_name = "pglib_opf_case300_ieee"
 formulation = "SOCWRConicPowerModel" # ACPPowerModel SOCWRConicPowerModel DCPPowerModel
 io_type = "input"
-download_dataset(organization, dataset, case_name, io_type; cache_dir=cache_dir)
+download_dataset(organization, dataset, case_name, io_type; cache_dir = cache_dir)
 
 io_type = "output"
-download_dataset(organization, dataset, case_name, io_type; formulation=formulation , cache_dir=cache_dir)
\ No newline at end of file
+download_dataset(
+    organization,
+    dataset,
+    case_name,
+    io_type;
+    formulation = formulation,
+    cache_dir = cache_dir,
+)
diff --git a/examples/powermodels/data_split.jl b/examples/powermodels/data_split.jl
index c02c320..159c98c 100644
--- a/examples/powermodels/data_split.jl
+++ b/examples/powermodels/data_split.jl
@@ -38,7 +38,8 @@ mkpath(joinpath(case_file_path_input, "test"))
 iter_files_in = readdir(joinpath(case_file_path_input))
 iter_files_in = filter(x -> occursin(string(filetype), x), iter_files_in)
 file_ins = [
-    joinpath(case_file_path_input, file) for file in iter_files_in if occursin("input", file)
+    joinpath(case_file_path_input, file) for
+    file in iter_files_in if occursin("input", file)
 ]
 batch_ids = [split(split(file, "_")[end], ".")[1] for file in file_ins]
 
@@ -56,7 +57,7 @@ input_data = DataFrame(input_table_train)
 # Split Data
 ##############
 Random.seed!(123)
-train_idx, test_idx = splitobs(1:size(input_data, 1), at=(0.7), shuffle=true)
+train_idx, test_idx = splitobs(1:size(input_data, 1), at = (0.7), shuffle = true)
 
 train_table = input_data[train_idx, :]
 test_table = input_data[test_idx, :]
@@ -72,10 +73,14 @@ num_batches = ceil(Int, length(test_idx) / batch_size)
 @info "Computing if test points are in the convex hull of the training set" batch_size num_batches
 
 inhull = Array{Bool}(undef, length(test_idx))
-@sync @distributed for i in 1:num_batches
-    idx_range = (i-1)*batch_size+1:min(i*batch_size, length(test_idx))
+@sync @distributed for i = 1:num_batches
+    idx_range = (i-1)*batch_size+1:min(i * batch_size, length(test_idx))
     batch = test_table[idx_range, :]
-    inhull[idx_range] = inconvexhull(Matrix(train_table[!, Not(:id)]), Matrix(batch[!, Not(:id)]), Gurobi.Optimizer)
+    inhull[idx_range] = inconvexhull(
+        Matrix(train_table[!, Not(:id)]),
+        Matrix(batch[!, Not(:id)]),
+        Gurobi.Optimizer,
+    )
     @info "Batch $i of $num_batches done"
 end
 
@@ -87,9 +92,21 @@ test_table.in_train_convex_hull = inhull
 
 # Save the training and test sets
 if filetype === ArrowFile
-    Arrow.write(joinpath(case_file_path_input, "train", case_name * "_train_input" * ".arrow"), train_table)
-    Arrow.write(joinpath(case_file_path_input, "test", case_name * "_test_input" * ".arrow"), test_table)
+    Arrow.write(
+        joinpath(case_file_path_input, "train", case_name * "_train_input" * ".arrow"),
+        train_table,
+    )
+    Arrow.write(
+        joinpath(case_file_path_input, "test", case_name * "_test_input" * ".arrow"),
+        test_table,
+    )
 else
-    CSV.write(joinpath(case_file_path_input, "train", case_name * "_train_input" * ".csv"), train_table)
-    CSV.write(joinpath(case_file_path_input, "test", case_name * "_test_input" * ".csv"), test_table)
-end
\ No newline at end of file
+    CSV.write(
+        joinpath(case_file_path_input, "train", case_name * "_train_input" * ".csv"),
+        train_table,
+    )
+    CSV.write(
+        joinpath(case_file_path_input, "test", case_name * "_test_input" * ".csv"),
+        test_table,
+    )
+end
diff --git a/examples/powermodels/generate_full_datasets_script.jl b/examples/powermodels/generate_full_datasets_script.jl
index d013bda..96a76bb 100644
--- a/examples/powermodels/generate_full_datasets_script.jl
+++ b/examples/powermodels/generate_full_datasets_script.jl
@@ -80,18 +80,25 @@ if haskey(config, "sampler")
     num_batches = config["sampler"]["num_batches"]
     num_p = config["sampler"]["num_samples"]
     global success_solves = 0.0
-    for i in 1:num_batches
-        _success_solves, number_variables, number_loads, batch_id = generate_dataset_pglib(
-            case_file_path,
-            case_name;
-            num_p=num_p,
-            filetype=filetype,
-            network_formulation=network_formulation,
-            optimizer=POI_cached_optimizer,
-            internal_load_sampler=(_o, n, idx, num_inputs) -> load_sampler(
-                _o, n, idx, num_inputs; max_multiplier=1.25, min_multiplier=0.8, step_multiplier=0.01
-            ),
-        )
+    for i = 1:num_batches
+        _success_solves, number_variables, number_loads, batch_id, _ =
+            generate_dataset_pglib(
+                case_file_path,
+                case_name;
+                num_p = num_p,
+                filetype = filetype,
+                network_formulation = network_formulation,
+                optimizer = POI_cached_optimizer,
+                internal_load_sampler = (_o, n, idx, num_inputs) -> load_sampler(
+                    _o,
+                    n,
+                    idx,
+                    num_inputs;
+                    max_multiplier = 1.25,
+                    min_multiplier = 0.8,
+                    step_multiplier = 0.01,
+                ),
+            )
         global success_solves += _success_solves
     end
     success_solves /= num_batches
@@ -110,18 +117,23 @@ if haskey(config, "line_search")
     early_stop_fn = (model, status, recorder) -> !status
 
     global success_solves = 0.0
-    for ibatc in 1:num_batches
-        _success_solves, number_variables, number_loads, b_id = generate_dataset_pglib(
+    for ibatc = 1:num_batches
+        _success_solves, number_variables, number_loads, b_id, _ = generate_dataset_pglib(
             case_file_path,
             case_name;
-            num_p=num_p,
-            filetype=filetype,
-            network_formulation=network_formulation,
-            optimizer=POI_cached_optimizer,
-            internal_load_sampler=(_o, n, idx, num_inputs) -> line_sampler(
-                _o, n, idx, num_inputs, ibatc; step_multiplier=step_multiplier
+            num_p = num_p,
+            filetype = filetype,
+            network_formulation = network_formulation,
+            optimizer = POI_cached_optimizer,
+            internal_load_sampler = (_o, n, idx, num_inputs) -> line_sampler(
+                _o,
+                n,
+                idx,
+                num_inputs,
+                ibatc;
+                step_multiplier = step_multiplier,
             ),
-            early_stop_fn=early_stop_fn,
+            early_stop_fn = early_stop_fn,
         )
         global success_solves += _success_solves
     end
diff --git a/examples/powermodels/jls2jld2.jl b/examples/powermodels/jls2jld2.jl
index 7296394..3bf3c28 100644
--- a/examples/powermodels/jls2jld2.jl
+++ b/examples/powermodels/jls2jld2.jl
@@ -39,12 +39,14 @@ end
 iter_files_in = readdir(joinpath(case_file_path_input))
 iter_files_in = filter(x -> occursin(string(filetype), x), iter_files_in)
 file_ins = [
-    joinpath(case_file_path_input, file) for file in iter_files_in if occursin("input", file)
+    joinpath(case_file_path_input, file) for
+    file in iter_files_in if occursin("input", file)
 ]
 iter_files_out = readdir(joinpath(case_file_path_output))
 iter_files_out = filter(x -> occursin(string(filetype), x), iter_files_out)
 file_outs = [
-    joinpath(case_file_path_output, file) for file in iter_files_out if occursin("output", file)
+    joinpath(case_file_path_output, file) for
+    file in iter_files_out if occursin("output", file)
 ]
 # batch_ids = [split(split(file, "_")[end], ".")[1] for file in file_ins]
 
@@ -62,10 +64,10 @@ input_data = DataFrame(input_table_train)
 output_data = DataFrame(output_table_train)
 
 # filter out rows with 0.0 operational_cost (i.e. inidicative of numerical issues)
-output_data = output_data[output_data.operational_cost .> 10, :]
+output_data = output_data[output_data.operational_cost.>10, :]
 
 # match
-train_table = innerjoin(input_data, output_data[!, [:id, :operational_cost]]; on=:id)
+train_table = innerjoin(input_data, output_data[!, [:id, :operational_cost]]; on = :id)
 
 input_features = names(train_table[!, Not([:id, :operational_cost])])
 
@@ -81,4 +83,9 @@ mach = machine(joinpath(model_dir, save_file * "$num.jls"))
 ##############
 model = mach.fitresult[1]
 model_state = Flux.state(model)
-jldsave(joinpath(model_dir, save_file * ".jld2"), model_state=model_state, input_features=input_features, layers=mach.model.builder.hidden_sizes)
\ No newline at end of file
+jldsave(
+    joinpath(model_dir, save_file * ".jld2"),
+    model_state = model_state,
+    input_features = input_features,
+    layers = mach.model.builder.hidden_sizes,
+)
diff --git a/examples/powermodels/pglib_datagen.jl b/examples/powermodels/pglib_datagen.jl
index 6fc1b25..1f504d1 100644
--- a/examples/powermodels/pglib_datagen.jl
+++ b/examples/powermodels/pglib_datagen.jl
@@ -30,9 +30,9 @@ function load_sampler(
     num_p::F,
     ::F,
     ::F;
-    max_multiplier::T=2.5,
-    min_multiplier::T=0.0,
-    step_multiplier::T=0.1,
+    max_multiplier::T = 2.5,
+    min_multiplier::T = 0.0,
+    step_multiplier::T = 0.1,
 ) where {T<:Real,F<:Integer}
     # Load sampling
     load_samples =
@@ -73,13 +73,16 @@ end
 
 load_parameter_factory is a function to help generate a parameter vector for the problem iterator.
 """
-function load_parameter_factory(model, indices; load_set=nothing)
+function load_parameter_factory(model, indices; load_set = nothing)
     if isnothing(load_set)
-        return @variable(model, _p[i=indices])
+        return @variable(model, _p[i = indices])
     end
-    num_loads = floor(Int,length(indices) / 2)
+    num_loads = floor(Int, length(indices) / 2)
     pd_index = indices[1:num_loads]
-    return [@variable(model, pd[i=pd_index] in load_set[i]).data; @variable(model, qd[i=pd_index] in load_set[i+num_loads]).data]
+    return [
+        @variable(model, pd[i = pd_index] in load_set[i]).data
+        @variable(model, qd[i = pd_index] in load_set[i+num_loads]).data
+    ]
 end
 
 """
@@ -92,14 +95,14 @@ function pm_primal_builder!(
     parameters,
     network_data,
     network_formulation;
-    recorder=nothing,
-    record_duals=false,
+    recorder = nothing,
+    record_duals = false,
 )
     num_loads = length(network_data["load"])
     for (str_i, l) in network_data["load"]
         i = parse(Int, str_i)
         l["pd"] = parameters[i]
-        l["qd"] = parameters[num_loads + i]
+        l["qd"] = parameters[num_loads+i]
     end
 
     # Instantiate the model
@@ -107,8 +110,8 @@ function pm_primal_builder!(
         network_data,
         network_formulation,
         PowerModels.build_opf;
-        setting=Dict("output" => Dict("branch_flows" => true, "duals" => true)),
-        jump_model=model,
+        setting = Dict("output" => Dict("branch_flows" => true, "duals" => true)),
+        jump_model = model,
     )
 
     if !isnothing(recorder)
@@ -120,10 +123,10 @@ function pm_primal_builder!(
         if record_duals
             # Bus
             real_balance = [bus[:lam_kcl_r] for bus in values(PowerModels.sol(pm)[:bus])]
-            set_name.(real_balance, ["lam_kcl_r[$(i)]" for i in 1:length(real_balance)])
+            set_name.(real_balance, ["lam_kcl_r[$(i)]" for i = 1:length(real_balance)])
             imag_balance = [bus[:lam_kcl_i] for bus in values(PowerModels.sol(pm)[:bus])]
             if eltype(imag_balance) <: ConstraintRef
-                set_name.(imag_balance, ["lam_kcl_i[$(i)]" for i in 1:length(imag_balance)])
+                set_name.(imag_balance, ["lam_kcl_i[$(i)]" for i = 1:length(imag_balance)])
                 set_dual_variable!(recorder, vcat(real_balance, imag_balance))
             else
                 set_dual_variable!(recorder, real_balance)
@@ -158,14 +161,14 @@ Generate dataset for pglib case_name with num_p problems and save it in data_dir
 function generate_dataset_pglib(
     data_dir,
     case_name;
-    filetype=CSVFile,
-    num_p=10,
-    internal_load_sampler=load_sampler,
-    network_formulation=DCPPowerModel,
-    optimizer=() -> POI.Optimizer(HiGHS.Optimizer()),
-    filterfn=LearningToOptimize.filter_fn,
-    early_stop_fn=(model, status, recorder) -> false,
-    batch_id=string(uuid1()),
+    filetype = CSVFile,
+    num_p = 10,
+    internal_load_sampler = load_sampler,
+    network_formulation = DCPPowerModel,
+    optimizer = () -> POI.Optimizer(HiGHS.Optimizer()),
+    filterfn = LearningToOptimize.filter_fn,
+    early_stop_fn = (model, status, recorder) -> false,
+    batch_id = string(uuid1()),
 )
     @info "Batch ID: $batch_id"
 
@@ -185,26 +188,36 @@ function generate_dataset_pglib(
     num_loads = length(network_data["load"])
     num_inputs = num_loads * 2
     original_load = vcat(
-        [network_data["load"]["$l"]["pd"] for l in 1:num_loads],
-        [network_data["load"]["$l"]["qd"] for l in 1:num_loads],
+        [network_data["load"]["$l"]["pd"] for l = 1:num_loads],
+        [network_data["load"]["$l"]["qd"] for l = 1:num_loads],
+    )
+    p = load_parameter_factory(
+        model,
+        1:num_inputs;
+        load_set = MOI.Parameter.(original_load),
     )
-    p = load_parameter_factory(model, 1:num_inputs; load_set=MOI.Parameter.(original_load))
 
     # Build model and Recorder
     file = joinpath(
-        data_sim_dir, case_name * "_" * string(network_formulation) * "_output_" * batch_id
+        data_sim_dir,
+        case_name * "_" * string(network_formulation) * "_output_" * batch_id,
     )
-    recorder = Recorder{filetype}(file; filterfn=filterfn,model=model)
+    recorder = Recorder{filetype}(file; filterfn = filterfn, model = model)
     pm_primal_builder!(
-        model, p, network_data, network_formulation; recorder=recorder, record_duals=true
+        model,
+        p,
+        network_data,
+        network_formulation;
+        recorder = recorder,
+        record_duals = true,
     )
 
     # The problem iterator
     pairs = Dict{VariableRef,Vector{Float64}}()
-    for i in 1:num_inputs
+    for i = 1:num_inputs
         pairs[p[i]] = internal_load_sampler(original_load[i], num_p, i, num_inputs)
     end
-    problem_iterator = ProblemIterator(pairs; early_stop=early_stop_fn)
+    problem_iterator = ProblemIterator(pairs; early_stop = early_stop_fn)
 
     save(
         problem_iterator,
@@ -219,18 +232,19 @@ function generate_dataset_pglib(
     return solve_batch(problem_iterator, recorder),
     length(recorder.primal_variables) + length(recorder.dual_variables),
     length(original_load),
-    batch_id
+    batch_id,
+    problem_iterator
 end
 
 function generate_worst_case_dataset_Nonconvex(
     data_dir,
     case_name;
-    filetype=CSVFile,
-    num_p=10,
-    network_formulation=DCPPowerModel,
-    optimizer=() -> POI.Optimizer(HiGHS.Optimizer()),
-    algorithm=NLoptAlg(:LN_BOBYQA),
-    options=NLoptOptions(; maxeval=10),
+    filetype = CSVFile,
+    num_p = 10,
+    network_formulation = DCPPowerModel,
+    optimizer = () -> POI.Optimizer(HiGHS.Optimizer()),
+    algorithm = NLoptAlg(:LN_BOBYQA),
+    options = NLoptOptions(; maxeval = 10),
 )
     # save folder
     data_sim_dir = joinpath(data_dir, string(network_formulation))
@@ -253,7 +267,9 @@ function generate_worst_case_dataset_Nonconvex(
         [l["qd"] for l in values(network_data["load"])],
     )
     p = load_parameter_factory(
-        model, 1:(num_loads * 2); load_set=MOI.Parameter.(original_load)
+        model,
+        1:(num_loads*2);
+        load_set = MOI.Parameter.(original_load),
     )
 
     # Define batch id
@@ -262,20 +278,25 @@ function generate_worst_case_dataset_Nonconvex(
 
     # File names
     file_input = joinpath(
-        data_sim_dir, case_name * "_" * string(network_formulation) * "_input_" * batch_id
+        data_sim_dir,
+        case_name * "_" * string(network_formulation) * "_input_" * batch_id,
     )
     file_output = joinpath(
-        data_sim_dir, case_name * "_" * string(network_formulation) * "_output_" * batch_id
+        data_sim_dir,
+        case_name * "_" * string(network_formulation) * "_output_" * batch_id,
     )
     recorder = Recorder{filetype}(
-        file_output; filename_input=file_input, primal_variables=[], dual_variables=[], model=model
+        file_output;
+        filename_input = file_input,
+        primal_variables = [],
+        dual_variables = [],
+        model = model,
     )
 
     # Build model
-    model, parameters, variable_refs = pm_primal_builder!(
-        model, p, network_data, network_formulation; recorder=recorder
-    )
-    function _primal_builder!(; recorder=nothing)
+    model, parameters, variable_refs =
+        pm_primal_builder!(model, p, network_data, network_formulation; recorder = recorder)
+    function _primal_builder!(; recorder = nothing)
         if !isnothing(recorder)
             set_primal_variable!(recorder, variable_refs)
         end
@@ -296,12 +317,12 @@ function generate_worst_case_dataset_Nonconvex(
 
     # The problem iterator
     problem_iterator = WorstCaseProblemIterator(
-        [uuid1() for _ in 1:num_p], # will be ignored
+        [uuid1() for _ = 1:num_p], # will be ignored
         () -> nothing, # will be ignored
         _primal_builder!,
         _set_iterator!,
         algorithm;
-        options=options,
+        options = options,
     )
 
     # Solve all problems and record solutions
@@ -318,11 +339,11 @@ end
 function generate_worst_case_dataset(
     data_dir,
     case_name;
-    filetype=CSVFile,
-    num_p=10,
-    network_formulation=DCPPowerModel,
-    optimizer_factory=default_optimizer_factory,
-    hook=nothing,
+    filetype = CSVFile,
+    num_p = 10,
+    network_formulation = DCPPowerModel,
+    optimizer_factory = default_optimizer_factory,
+    hook = nothing,
 )
     # save folder
     data_sim_dir = joinpath(data_dir, string(network_formulation))
@@ -340,7 +361,7 @@ function generate_worst_case_dataset(
         [l["pd"] for l in values(network_data["load"])],
         [l["qd"] for l in values(network_data["load"])],
     )
-    parameter_factory = (model) -> load_parameter_factory(model, 1:(num_loads * 2))
+    parameter_factory = (model) -> load_parameter_factory(model, 1:(num_loads*2))
 
     # Define batch id
     batch_id = string(uuid1())
@@ -348,8 +369,12 @@ function generate_worst_case_dataset(
 
     # Build model
     primal_builder! =
-        (model, parameters; recorder=nothing) -> pm_primal_builder!(
-            model, parameters, network_data, network_formulation; recorder=recorder
+        (model, parameters; recorder = nothing) -> pm_primal_builder!(
+            model,
+            parameters,
+            network_data,
+            network_formulation;
+            recorder = recorder,
         )
 
     # Set iterator
@@ -359,23 +384,29 @@ function generate_worst_case_dataset(
 
     # The problem iterator
     problem_iterator = WorstCaseProblemIterator(
-        [uuid1() for _ in 1:num_p],
+        [uuid1() for _ = 1:num_p],
         parameter_factory,
         primal_builder!,
         set_iterator!,
         optimizer_factory;
-        hook=hook,
+        hook = hook,
     )
 
     # File names
     file_input = joinpath(
-        data_sim_dir, case_name * "_" * string(network_formulation) * "_input_" * batch_id
+        data_sim_dir,
+        case_name * "_" * string(network_formulation) * "_input_" * batch_id,
     )
     file_output = joinpath(
-        data_sim_dir, case_name * "_" * string(network_formulation) * "_output_" * batch_id
+        data_sim_dir,
+        case_name * "_" * string(network_formulation) * "_output_" * batch_id,
     )
     recorder = Recorder{filetype}(
-        file_output; filename_input=file_input, primal_variables=[], dual_variables=[], model=JuMP.Model() # dummy model
+        file_output;
+        filename_input = file_input,
+        primal_variables = [],
+        dual_variables = [],
+        model = JuMP.Model(), # dummy model
     )
 
     # Solve all problems and record solutions
@@ -388,9 +419,13 @@ end
 function test_pglib_datasetgen(path::AbstractString, case_name::AbstractString, num_p::Int)
     @testset "Dataset Generation pglib case" begin
         network_formulation = DCPPowerModel
-        success_solves, number_variables, number_parameters, batch_id = generate_dataset_pglib(
-            path, case_name; num_p=num_p, network_formulation=network_formulation
-        )
+        success_solves, number_variables, number_parameters, batch_id, problem_iterator =
+            generate_dataset_pglib(
+                path,
+                case_name;
+                num_p = num_p,
+                network_formulation = network_formulation,
+            )
         file_in = joinpath(
             path,
             string(network_formulation),
@@ -411,19 +446,25 @@ function test_pglib_datasetgen(path::AbstractString, case_name::AbstractString,
         @test length(readdlm(file_out, ',')[:, 1]) == num_p * success_solves + 1
         @test length(readdlm(file_out, ',')[1, :]) == number_variables + 6
 
-        return file_in, file_out
+        return file_in, file_out, problem_iterator
     end
 end
 
 function test_generate_worst_case_dataset(
-    path::AbstractString, case_name::AbstractString, num_p::Int
+    path::AbstractString,
+    case_name::AbstractString,
+    num_p::Int,
 )
     @testset "Worst Case Dataset Generation pglib case" begin
         network_formulation = DCPPowerModel
         # Improve dataset
-        success_solves, number_variables, number_parameters, batch_id = generate_worst_case_dataset(
-            path, case_name; num_p=num_p, network_formulation=network_formulation
-        )
+        success_solves, number_variables, number_parameters, batch_id =
+            generate_worst_case_dataset(
+                path,
+                case_name;
+                num_p = num_p,
+                network_formulation = network_formulation,
+            )
 
         file_in = joinpath(
             path,
@@ -451,14 +492,20 @@ function test_generate_worst_case_dataset(
 end
 
 function test_generate_worst_case_dataset_Nonconvex(
-    path::AbstractString, case_name::AbstractString, num_p::Int
+    path::AbstractString,
+    case_name::AbstractString,
+    num_p::Int,
 )
     @testset "WC Nonconvex Dataset Generation pglib case" begin
         network_formulation = DCPPowerModel
         # Improve dataset
-        success_solves, number_variables, number_parameters, batch_id = generate_worst_case_dataset_Nonconvex(
-            path, case_name; num_p=num_p, network_formulation=network_formulation
-        )
+        success_solves, number_variables, number_parameters, batch_id =
+            generate_worst_case_dataset_Nonconvex(
+                path,
+                case_name;
+                num_p = num_p,
+                network_formulation = network_formulation,
+            )
 
         file_in = joinpath(
             path,
diff --git a/examples/powermodels/powermodels_nn_testing.jl b/examples/powermodels/powermodels_nn_testing.jl
index 6e8daf6..dd1edd6 100644
--- a/examples/powermodels/powermodels_nn_testing.jl
+++ b/examples/powermodels/powermodels_nn_testing.jl
@@ -30,12 +30,14 @@ save_file = "$(case_name)_$(network_formulation)"
 iter_files_in = readdir(joinpath(case_file_path_input))
 iter_files_in = filter(x -> occursin(string(filetype), x), iter_files_in)
 file_ins = [
-    joinpath(case_file_path_input, file) for file in iter_files_in if occursin("input", file)
+    joinpath(case_file_path_input, file) for
+    file in iter_files_in if occursin("input", file)
 ]
 iter_files_out = readdir(joinpath(case_file_path_output))
 iter_files_out = filter(x -> occursin(string(filetype), x), iter_files_out)
 file_outs = [
-    joinpath(case_file_path_output, file) for file in iter_files_out if occursin("output", file)
+    joinpath(case_file_path_output, file) for
+    file in iter_files_out if occursin("output", file)
 ]
 
 # Load input and output data tables
@@ -52,10 +54,10 @@ input_data = DataFrame(input_table_train)
 output_data = DataFrame(output_table_train)
 
 # filter out rows with 0.0 operational_cost (i.e. inidicative of numerical issues)
-output_data = output_data[output_data.operational_cost .> 10, :]
+output_data = output_data[output_data.operational_cost.>10, :]
 
 # match
-test_table = innerjoin(input_data, output_data[!, [:id, :operational_cost]]; on=:id)
+test_table = innerjoin(input_data, output_data[!, [:id, :operational_cost]]; on = :id)
 
 ##############
 # Load DNN Approximator
@@ -102,8 +104,8 @@ std_out_convex_hull = Array{Float32}(undef, length(flux_models))
 
 for (i, flux_model) in enumerate(flux_models)
     error_vec = (y .- flux_model(X')') ./ y
-    error_vec_in_chull = error_vec[test_table.in_train_convex_hull .== 1]
-    error_vec_out_chull = error_vec[test_table.in_train_convex_hull .== 0]
+    error_vec_in_chull = error_vec[test_table.in_train_convex_hull.==1]
+    error_vec_out_chull = error_vec[test_table.in_train_convex_hull.==0]
     mae_convex_hull[i] = mean(abs.(error_vec_in_chull))
     mae_out_convex_hull[i] = mean(abs.(error_vec_out_chull))
     worst_case_convex_hull[i] = maximum(abs.(error_vec_in_chull))
@@ -129,4 +131,4 @@ results_dir = joinpath(dirname(@__FILE__), "results")
 mkpath(results_dir)
 save_file = joinpath(results_dir, "$(save_file)_results.csv")
 
-CSV.write(save_file, results)
\ No newline at end of file
+CSV.write(save_file, results)
diff --git a/examples/powermodels/powermodels_nn_training.jl b/examples/powermodels/powermodels_nn_training.jl
index 8cb58eb..fb79b40 100644
--- a/examples/powermodels/powermodels_nn_training.jl
+++ b/examples/powermodels/powermodels_nn_training.jl
@@ -43,12 +43,14 @@ end
 iter_files_in = readdir(joinpath(case_file_path_input))
 iter_files_in = filter(x -> occursin(string(filetype), x), iter_files_in)
 file_ins = [
-    joinpath(case_file_path_input, file) for file in iter_files_in if occursin("input", file)
+    joinpath(case_file_path_input, file) for
+    file in iter_files_in if occursin("input", file)
 ]
 iter_files_out = readdir(joinpath(case_file_path_output))
 iter_files_out = filter(x -> occursin(string(filetype), x), iter_files_out)
 file_outs = [
-    joinpath(case_file_path_output, file) for file in iter_files_out if occursin("output", file)
+    joinpath(case_file_path_output, file) for
+    file in iter_files_out if occursin("output", file)
 ]
 # batch_ids = [split(split(file, "_")[end], ".")[1] for file in file_ins]
 
@@ -66,10 +68,10 @@ input_data = DataFrame(input_table_train)
 output_data = DataFrame(output_table_train)
 
 # filter out rows with 0.0 operational_cost (i.e. inidicative of numerical issues)
-output_data = output_data[output_data.operational_cost .> 10, :]
+output_data = output_data[output_data.operational_cost.>10, :]
 
 # match
-train_table = innerjoin(input_data, output_data[!, [:id, :operational_cost]]; on=:id)
+train_table = innerjoin(input_data, output_data[!, [:id, :operational_cost]]; on = :id)
 
 input_features = names(train_table[!, Not([:id, :operational_cost])])
 
@@ -92,13 +94,23 @@ lg = WandbLogger(
         "rng" => 123,
         "network_formulation" => network_formulation,
         # "lambda" => 0.00,
-    )
+    ),
 )
 
-optimiser= Flux.Optimise.Adam(get_config(lg, "learning_rate"), (0.9, 0.999), 1.0e-8, IdDict{Any,Any}())
+optimiser = Flux.Optimise.Adam(
+    get_config(lg, "learning_rate"),
+    (0.9, 0.999),
+    1.0e-8,
+    IdDict{Any,Any}(),
+)
 if icnn
-    optimiser=ConvexRule(
-        Flux.Optimise.Adam(get_config(lg, "learning_rate"), (0.9, 0.999), 1.0e-8, IdDict{Any,Any}())
+    optimiser = ConvexRule(
+        Flux.Optimise.Adam(
+            get_config(lg, "learning_rate"),
+            (0.9, 0.999),
+            1.0e-8,
+            IdDict{Any,Any}(),
+        ),
     )
 end
 
@@ -114,14 +126,14 @@ end
 # )
 
 nn = MultitargetNeuralNetworkRegressor(;
-    builder=FullyConnectedBuilder(layers),
-    rng=get_config(lg, "rng"),
-    epochs=5000,
-    optimiser=optimiser,
-    acceleration=CUDALibs(),
-    batch_size=get_config(lg, "batch_size"),
+    builder = FullyConnectedBuilder(layers),
+    rng = get_config(lg, "rng"),
+    epochs = 5000,
+    optimiser = optimiser,
+    acceleration = CUDALibs(),
+    batch_size = get_config(lg, "batch_size"),
     # lambda=get_config(lg, "lambda"),
-    loss=relative_rmse,
+    loss = relative_rmse,
 )
 
 # Constrols
@@ -136,30 +148,31 @@ model_path = joinpath(model_dir, save_file * ".jld2")
 
 save_control = SaveBest(1000, model_path, 0.003)
 
-controls=[Step(1),
-    WithModelLossDo(save_control; stop_if_true=true),
+controls = [
+    Step(1),
+    WithModelLossDo(save_control; stop_if_true = true),
     # NumberLimit(n=4),
     # NumberSinceBest(6),
     # PQ(; alpha=0.9, k=30),
     # GL(; alpha=4.0),
     InvalidValue(),
     # Threshold(0.003),
-    TimeLimit(; t=3),
+    TimeLimit(; t = 3),
     WithLossDo(update_loss),
     WithReportDo(update_training_loss),
-    WithIterationsDo(update_epochs)
+    WithIterationsDo(update_epochs),
 ]
 
-iterated_pipe =
-    IteratedModel(model=nn,
-        controls=controls,
-        # resampling=Holdout(fraction_train=0.7),
-        measure = relative_mae,
+iterated_pipe = IteratedModel(
+    model = nn,
+    controls = controls,
+    # resampling=Holdout(fraction_train=0.7),
+    measure = relative_mae,
 )
 
 # Fit model
 mach = machine(iterated_pipe, X, y)
-fit!(mach; verbosity=2)
+fit!(mach; verbosity = 2)
 
 # Finish the run
 close(lg)
@@ -173,4 +186,4 @@ close(lg)
 #     model_state = Flux.state(fitted_model)
 
 #     jldsave(model_path; model_state=model_state, layers=layers, input_features=input_features)
-# end
\ No newline at end of file
+# end
diff --git a/examples/powermodels/visualize.jl b/examples/powermodels/visualize.jl
index 8f8b53c..314ddb9 100644
--- a/examples/powermodels/visualize.jl
+++ b/examples/powermodels/visualize.jl
@@ -4,7 +4,7 @@ using DataFrames
 using Statistics
 using LinearAlgebra
 
-cossim(x,y) = dot(x,y) / (norm(x)*norm(y))
+cossim(x, y) = dot(x, y) / (norm(x) * norm(y))
 
 ##############
 # Parameters
@@ -25,15 +25,18 @@ iter_files_test = readdir(joinpath(case_file_path_test))
 iter_files_train = filter(x -> occursin("arrow", x), iter_files_train)
 iter_files_test = filter(x -> occursin("arrow", x), iter_files_test)
 file_train = [
-    joinpath(case_file_path_train, file) for file in iter_files_train if occursin("input", file)
+    joinpath(case_file_path_train, file) for
+    file in iter_files_train if occursin("input", file)
 ]
 file_test = [
-    joinpath(case_file_path_test, file) for file in iter_files_test if occursin("input", file)
+    joinpath(case_file_path_test, file) for
+    file in iter_files_test if occursin("input", file)
 ]
 iter_files_out = readdir(joinpath(case_file_path_output))
 iter_files_out = filter(x -> occursin("arrow", x), iter_files_out)
 file_outs = [
-    joinpath(case_file_path_output, file) for file in iter_files_out if occursin("output", file)
+    joinpath(case_file_path_output, file) for
+    file in iter_files_out if occursin("output", file)
 ]
 
 # Load input and output data tables
@@ -46,28 +49,37 @@ input_data_train = DataFrame(input_table_train)
 input_data_test = DataFrame(input_table_test)
 output_data = DataFrame(output_table)
 output_data.operational_cost = output_data.objective
-output_data = output_data[output_data.objective .> 10, :]
+output_data = output_data[output_data.objective.>10, :]
 input_data = vcat(input_data_train, input_data_test[!, Not(:in_train_convex_hull)])
 
 ##############
 # SOC VS AC
 ##############
 network_formulation_soc = "SOCWRConicPowerModel"
-case_file_path_output_soc = joinpath(case_file_path, "output", string(network_formulation_soc))
+case_file_path_output_soc =
+    joinpath(case_file_path, "output", string(network_formulation_soc))
 iter_files_out_soc = readdir(case_file_path_output_soc)
 iter_files_out_soc = filter(x -> occursin("arrow", x), iter_files_out_soc)
 file_outs_soc = [
-    joinpath(case_file_path_output_soc, file) for file in iter_files_out_soc if occursin("output", file)
+    joinpath(case_file_path_output_soc, file) for
+    file in iter_files_out_soc if occursin("output", file)
 ]
 output_table_soc = Arrow.Table(file_outs_soc)
 output_data_soc = DataFrame(output_table_soc)
 output_data_soc.operational_cost_soc = output_data_soc.objective
-output_data_soc = output_data_soc[output_data_soc.objective .> 10, :]
+output_data_soc = output_data_soc[output_data_soc.objective.>10, :]
 
 # compare SOC and AC operational_cost by id
-ac_soc = innerjoin(output_data[!, [:id, :operational_cost]], output_data_soc[!, [:id, :operational_cost_soc]], on=:id, makeunique=true)
-
-ac_soc.error = (ac_soc.operational_cost .- ac_soc.operational_cost_soc) ./ ac_soc.operational_cost * 100
+ac_soc = innerjoin(
+    output_data[!, [:id, :operational_cost]],
+    output_data_soc[!, [:id, :operational_cost_soc]],
+    on = :id,
+    makeunique = true,
+)
+
+ac_soc.error =
+    (ac_soc.operational_cost .- ac_soc.operational_cost_soc) ./ ac_soc.operational_cost *
+    100
 mean(ac_soc.error)
 maximum(ac_soc.error)
 ac_soc[findmax(ac_soc.error)[2], :]
@@ -76,17 +88,17 @@ ac_soc[findmax(ac_soc.error)[2], :]
 ##############
 
 # Load vectors
-function total_load_vector(input_data; is_test=false)
+function total_load_vector(input_data; is_test = false)
     df = DataFrame()
     df.id = input_data.id
     if is_test
         df.in_hull = input_data.in_train_convex_hull
     end
     num_loads = length([col for col in names(input_data) if occursin("pd", col)])
-    for i in 1:num_loads
+    for i = 1:num_loads
         df[!, "load[$i]"] = zeros(size(input_data, 1))
     end
-    for j in 1:size(input_data, 1), i in 1:num_loads
+    for j = 1:size(input_data, 1), i = 1:num_loads
         df[j, "load[$i]"] = sqrt(input_data[j, "pd[$i]"]^2 + input_data[j, "qd[$i]"]^2)
     end
     return df
@@ -96,10 +108,10 @@ function total_load_vector_annon(input_data)
     df = DataFrame()
     df.id = input_data.id
     num_loads = floor(Int, length(names(input_data[!, Not(:id)])) / 2)
-    for i in 1:num_loads
+    for i = 1:num_loads
         df[!, "load[$i]"] = zeros(size(input_data, 1))
     end
-    for j in 1:size(input_data, 1), i in 1:num_loads
+    for j = 1:size(input_data, 1), i = 1:num_loads
         df[j, "load[$i]"] = sqrt(input_data[j, i+1]^2 + input_data[j, i+num_loads+1]^2)
     end
     return df
@@ -107,25 +119,38 @@ end
 
 ######### Plot Load Vectors #########
 load_vector_train = total_load_vector(input_data_train)
-load_vector_test = total_load_vector(input_data_test; is_test=true)
+load_vector_test = total_load_vector(input_data_test; is_test = true)
 
 # Nominal Loads
 nominal_loads = Vector(input_data_train[1, Not(:id)])
 norm_nominal_loads = norm(nominal_loads)
 
 # Load divergence
-theta_train = [acos(cossim(nominal_loads, Vector(input_data_train[i, Not(:id)]))) for i in 2:10000] * 180 / pi
-norm_sim_train = [norm(Vector(input_data_train[i, Not(:id)])) / norm_nominal_loads for i in 2:10000]
-
-theta_test = [acos(cossim(nominal_loads, Vector(load_vector_test[i, Not([:id, :in_hull])]))) for i in 1:size(load_vector_test, 1)] * 180 / pi
-norm_sim_test = [norm(Vector(load_vector_test[i, Not([:id, :in_hull])])) / norm_nominal_loads for i in 1:size(load_vector_test, 1)]
+theta_train =
+    [acos(cossim(nominal_loads, Vector(input_data_train[i, Not(:id)]))) for i = 2:10000] * 180 /
+    pi
+norm_sim_train =
+    [norm(Vector(input_data_train[i, Not(:id)])) / norm_nominal_loads for i = 2:10000]
+
+theta_test =
+    [
+        acos(cossim(nominal_loads, Vector(load_vector_test[i, Not([:id, :in_hull])]))) for
+        i = 1:size(load_vector_test, 1)
+    ] * 180 / pi
+norm_sim_test = [
+    norm(Vector(load_vector_test[i, Not([:id, :in_hull])])) / norm_nominal_loads for
+    i = 1:size(load_vector_test, 1)
+]
 
 # Polar Plot divergence
 gr()
 l = @layout [a b]
-p1 = plot(scatter(theta_train, norm_sim_train, proj = :polar, label="Train", color=:blue));
-p2 = plot(scatter(theta_test, norm_sim_test, proj = :polar, label="Test", color=:orange));
-plot(p1, p2; title="Load Similarity", legend=:bottomright, layout = l)
+p1 = plot(
+    scatter(theta_train, norm_sim_train, proj = :polar, label = "Train", color = :blue),
+);
+p2 =
+    plot(scatter(theta_test, norm_sim_test, proj = :polar, label = "Test", color = :orange));
+plot(p1, p2; title = "Load Similarity", legend = :bottomright, layout = l)
 
 ######### Plot Objective Function #########
 # Plot objective function for a single direction
@@ -136,19 +161,20 @@ load_vector = total_load_vector(input_data)
 nominal_loads = Vector(load_vector[1, Not(:id)])
 norm_nominal_loads = norm(nominal_loads)
 
-load_vector.norm_loads = [norm(Vector(load_vector[i, Not(:id)])) for i in 1:size(load_vector, 1)] ./ norm_nominal_loads
+load_vector.norm_loads =
+    [norm(Vector(load_vector[i, Not(:id)])) for i = 1:size(load_vector, 1)] ./ norm_nominal_loads
 # join input and output data
-joined_data = innerjoin(load_vector, output_data[!, [:id, :operational_cost]], on=:id)
+joined_data = innerjoin(load_vector, output_data[!, [:id, :operational_cost]], on = :id)
 
 # get k extreme points
 using LearningToOptimize
 using Gurobi
 function maxk(a, k)
-    b = partialsortperm(a, 1:k, rev=true)
+    b = partialsortperm(a, 1:k, rev = true)
     return collect(zip(b, a[b]))
 end
 function mink(a, k)
-    b = partialsortperm(a, 1:k, rev=false)
+    b = partialsortperm(a, 1:k, rev = false)
     return collect(zip(b, a[b]))
 end
 k = 10
@@ -157,26 +183,46 @@ idx_mins = mink(joined_data.norm_loads, k)
 
 # Objective function
 instances_in_convex_hulls = Array{DataFrame}(undef, k)
-for i in 1:k
+for i = 1:k
     @info "Processing instance $i"
     instance_load_max = joined_data[idx_maxs[i][1]:idx_maxs[i][1], :]
     instance_load_min = joined_data[idx_mins[i][1]:idx_mins[i][1], :]
 
     # find instances between the two points (i.e. in their convex hull)
     in_convex_hull = inconvexhull(
-        Matrix(vcat(instance_load_max, 
-        instance_load_min)[!, Not([:id, :norm_loads, :operational_cost])]), Matrix(joined_data[!, Not([:id, :norm_loads, :operational_cost])]),
-        Gurobi.Optimizer, silent=false, tol=0.1 # close to convex hull
+        Matrix(
+            vcat(instance_load_max, instance_load_min)[
+                !,
+                Not([:id, :norm_loads, :operational_cost]),
+            ],
+        ),
+        Matrix(joined_data[!, Not([:id, :norm_loads, :operational_cost])]),
+        Gurobi.Optimizer,
+        silent = false,
+        tol = 0.1, # close to convex hull
     )
 
-    instances_in_convex_hulls[i] = sort(joined_data[in_convex_hull, [:operational_cost, :norm_loads]], :norm_loads)
+    instances_in_convex_hulls[i] =
+        sort(joined_data[in_convex_hull, [:operational_cost, :norm_loads]], :norm_loads)
 end
 
 # Plot
 plotly()
 
-plt = plot(color=:blue, legend=false, xlabel="Total Load (%)", ylabel="Operational Cost", title="AC OPF - IEEE300");
-for i in 1:k
-    plot!(instances_in_convex_hulls[i][!, :norm_loads] * 100, instances_in_convex_hulls[i][!, :operational_cost], label="Direction $i", color=:red, alpha=0.4)
+plt = plot(
+    color = :blue,
+    legend = false,
+    xlabel = "Total Load (%)",
+    ylabel = "Operational Cost",
+    title = "AC OPF - IEEE300",
+);
+for i = 1:k
+    plot!(
+        instances_in_convex_hulls[i][!, :norm_loads] * 100,
+        instances_in_convex_hulls[i][!, :operational_cost],
+        label = "Direction $i",
+        color = :red,
+        alpha = 0.4,
+    )
 end
 plt
diff --git a/examples/powermodels/write_to_file.jl b/examples/powermodels/write_to_file.jl
index 3627a46..60ca0fb 100644
--- a/examples/powermodels/write_to_file.jl
+++ b/examples/powermodels/write_to_file.jl
@@ -12,20 +12,20 @@ cached_clara =
             Clarabel.Optimizer(),
         ),
         Float64,
-)
+    )
 
 POI_cached_optimizer_clara() = POI.Optimizer(cached_clara())
 
 ipopt = Ipopt.Optimizer()
-    MOI.set(ipopt, MOI.RawOptimizerAttribute("print_level"), 0)
-    cached =
-        () -> MOI.Bridges.full_bridge_optimizer(
-            MOI.Utilities.CachingOptimizer(
-                MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
-                ipopt,
-            ),
-            Float64,
-)
+MOI.set(ipopt, MOI.RawOptimizerAttribute("print_level"), 0)
+cached =
+    () -> MOI.Bridges.full_bridge_optimizer(
+        MOI.Utilities.CachingOptimizer(
+            MOI.Utilities.UniversalFallback(MOI.Utilities.Model{Float64}()),
+            ipopt,
+        ),
+        Float64,
+    )
 POI_cached_optimizer() = POI.Optimizer(cached())
 
 network_formulation = ACPPowerModel # ACPPowerModel SOCWRConicPowerModel DCPPowerModel
@@ -40,7 +40,7 @@ model = JuMP.Model(() -> POI_cached_optimizer())
 # Save original load value and Link POI
 num_loads = length(network_data["load"])
 
-@variable(model, load_scaler[i=1:num_loads] in MOI.Parameter.(1.0)) # needs fixing -> need to be instantiated after all other variables
+@variable(model, load_scaler[i = 1:num_loads] in MOI.Parameter.(1.0)) # needs fixing -> need to be instantiated after all other variables
 
 for (str_i, l) in network_data["load"]
     i = parse(Int, str_i)
@@ -52,8 +52,8 @@ pm = instantiate_model(
     network_data,
     network_formulation,
     PowerModels.build_opf;
-    setting=Dict("output" => Dict("branch_flows" => true, "duals" => true)),
-    jump_model=model,
+    setting = Dict("output" => Dict("branch_flows" => true, "duals" => true)),
+    jump_model = model,
 )
 
 JuMP.optimize!(model)
@@ -66,4 +66,4 @@ write_to_file(model, "$(matpower_case_name)_$(network_formulation)_POI_load.mof.
 # set_optimizer(dest_model, () -> POI_cached_optimizer())
 # JuMP.optimize!(dest_model)
 # JuMP.termination_status(dest_model)
-# JuMP.objective_value(dest_model)
\ No newline at end of file
+# JuMP.objective_value(dest_model)
diff --git a/examples/slurm/slurm.jl b/examples/slurm/slurm.jl
index 3839463..522879d 100644
--- a/examples/slurm/slurm.jl
+++ b/examples/slurm/slurm.jl
@@ -7,10 +7,10 @@ If you get an error make sure to Pkg.checkout("CluterManagers").
 
 try
 
-	using Distributed, ClusterManagers
+    using Distributed, ClusterManagers
 catch
-	Pkg.add("ClusterManagers")
-	Pkg.checkout("ClusterManagers")
+    Pkg.add("ClusterManagers")
+    Pkg.checkout("ClusterManagers")
 end
 
 using Distributed, ClusterManagers
diff --git a/examples/solve_dataset.jl b/examples/solve_dataset.jl
index a03e07e..163d459 100644
--- a/examples/solve_dataset.jl
+++ b/examples/solve_dataset.jl
@@ -38,12 +38,19 @@ using Random
 
 
 ########## PARAMETERS ##########
-model_file = joinpath(l2o_path, "examples/powermodels/data/6468_rte/input/6468_rte_ACPPowerModel_POI_load.mof.json") # ACPPowerModel SOCWRConicPowerModel DCPPowerModel
-input_file = joinpath(l2o_path, "examples/powermodels/data/6468_rte/input/6468_rte_POI_load_input_7f284054-d107-11ee-3fe9-09f5e129b1ad")
+model_file = joinpath(
+    l2o_path,
+    "examples/powermodels/data/6468_rte/input/6468_rte_ACPPowerModel_POI_load.mof.json",
+) # ACPPowerModel SOCWRConicPowerModel DCPPowerModel
+input_file = joinpath(
+    l2o_path,
+    "examples/powermodels/data/6468_rte/input/6468_rte_POI_load_input_7f284054-d107-11ee-3fe9-09f5e129b1ad",
+)
 
 save_path = joinpath(l2o_path, "examples/powermodels/data/6468_rte/output/ACPPowerModel")
 case_name = split(split(model_file, ".mof.")[1], "/")[end]
-processed_output_files = [file for file in readdir(save_path; join=true) if occursin(case_name, file)]
+processed_output_files =
+    [file for file in readdir(save_path; join = true) if occursin(case_name, file)]
 ids = if length(processed_output_files) == 0
     UUID[]
 else
@@ -53,9 +60,10 @@ batch_size = 200
 
 ########## SOLVE ##########
 
-problem_iterator_factory, num_batches = load(model_file, input_file, filetype; batch_size=batch_size, ignore_ids=ids)
+problem_iterator_factory, num_batches =
+    load(model_file, input_file, filetype; batch_size = batch_size, ignore_ids = ids)
 
-@sync @distributed for i in 1:num_batches
+@sync @distributed for i = 1:num_batches
     ipopt = Ipopt.Optimizer()
     MOI.set(ipopt, MOI.RawOptimizerAttribute("print_level"), 0)
     cached =
@@ -65,14 +73,24 @@ problem_iterator_factory, num_batches = load(model_file, input_file, filetype; b
                 ipopt,
             ),
             Float64,
-    )
+        )
     POI_cached_optimizer() = POI.Optimizer(cached())
     batch_id = uuid1()
     problem_iterator = problem_iterator_factory(i)
     set_optimizer(problem_iterator.model, () -> POI_cached_optimizer())
     output_file = joinpath(save_path, "$(case_name)_output_$(batch_id)")
-    recorder = Recorder{filetype}(output_file; filterfn= (model) -> true, model=problem_iterator.model)
+    recorder = Recorder{filetype}(
+        output_file;
+        filterfn = (model) -> true,
+        model = problem_iterator.model,
+    )
     successfull_solves = solve_batch(problem_iterator, recorder)
     @info "Solved $(length(successfull_solves)) problems"
-    LearningToOptimize.compress_batch_arrow(save_path, case_name; keyword_all="output", batch_id=string(batch_id), keyword_any=[string(batch_id)])
+    LearningToOptimize.compress_batch_arrow(
+        save_path,
+        case_name;
+        keyword_all = "output",
+        batch_id = string(batch_id),
+        keyword_any = [string(batch_id)],
+    )
 end
diff --git a/examples/training_utils.jl b/examples/training_utils.jl
index d38e4ff..302b0f9 100644
--- a/examples/training_utils.jl
+++ b/examples/training_utils.jl
@@ -11,7 +11,7 @@ function update_training_loss(report)
 end
 
 function update_epochs(epoch)
-    Wandb.log(lg, Dict("log/epoch" => epoch)) 
+    Wandb.log(lg, Dict("log/epoch" => epoch))
     return nothing
 end
 
@@ -24,23 +24,27 @@ struct WithModelLossDo{F<:Function}
 end
 
 # constructor:
-WithModelLossDo(; f=(x,model)->@info("loss: $x"),
-           stop_if_true=false,
-           stop_message=nothing) = WithModelLossDo(f, stop_if_true, stop_message)
-WithModelLossDo(f; kwargs...) = WithModelLossDo(; f=f, kwargs...)
+WithModelLossDo(;
+    f = (x, model) -> @info("loss: $x"),
+    stop_if_true = false,
+    stop_message = nothing,
+) = WithModelLossDo(f, stop_if_true, stop_message)
+WithModelLossDo(f; kwargs...) = WithModelLossDo(; f = f, kwargs...)
 
 IterationControl.needs_loss(::WithModelLossDo) = true
 
-function IterationControl.update!(c::WithModelLossDo,
-                 model,
-                 verbosity,
-                 n,
-                 state=(loss=nothing, done=false))
+function IterationControl.update!(
+    c::WithModelLossDo,
+    model,
+    verbosity,
+    n,
+    state = (loss = nothing, done = false),
+)
     loss = IterationControl.loss(model)
     loss === nothing && throw(ERR_NEEDS_LOSS)
     r = c.f(loss, model)
     done = (c.stop_if_true && r isa Bool && r) ? true : false
-    return (loss=loss, done=done)
+    return (loss = loss, done = done)
 end
 
 IterationControl.done(c::WithModelLossDo, state) = state.done
@@ -48,8 +52,8 @@ IterationControl.done(c::WithModelLossDo, state) = state.done
 function IterationControl.takedown(c::WithModelLossDo, verbosity, state)
     verbosity > 1 && @info "final loss: $(state.loss). "
     if state.done
-        message = c.stop_message === nothing ?
-            "Stop triggered by a `WithModelLossDo` control. " :
+        message =
+            c.stop_message === nothing ? "Stop triggered by a `WithModelLossDo` control. " :
             c.stop_message
         verbosity > 0 && @info message
         return merge(state, (log = message,))
@@ -72,7 +76,12 @@ function (callback::SaveBest)(loss, ic_model)
         callback.best_loss = loss
         model = mach.fitresult[1]
         model_state = Flux.state(model)
-        jldsave(model_path; model_state=model_state, layers=layers, input_features=input_features)
+        jldsave(
+            model_path;
+            model_state = model_state,
+            layers = layers,
+            input_features = input_features,
+        )
     end
     if loss < callback.threshold
         return true
diff --git a/examples/unitcommitment/bnb_dataset.jl b/examples/unitcommitment/bnb_dataset.jl
index 4138438..dd7fac6 100644
--- a/examples/unitcommitment/bnb_dataset.jl
+++ b/examples/unitcommitment/bnb_dataset.jl
@@ -18,14 +18,15 @@ using SparseArrays
 using Statistics
 using Random
 
-import UnitCommitment:
-    Formulation,
-    KnuOstWat2018,
-    MorLatRam2013,
-    ShiftFactorsFormulation
+import UnitCommitment: Formulation, KnuOstWat2018, MorLatRam2013, ShiftFactorsFormulation
 
-function uc_load_disturbances!(rng, instance, nominal_loads; load_disturbances_range=-20:20)
-    for i in 1:length(instance.buses)
+function uc_load_disturbances!(
+    rng,
+    instance,
+    nominal_loads;
+    load_disturbances_range = -20:20,
+)
+    for i = 1:length(instance.buses)
         bus = instance.buses[i]
         bus.load = nominal_loads[i] .+ rand(rng, load_disturbances_range)
         bus.load = max.(bus.load, 0.0)
@@ -35,12 +36,14 @@ end
 """
     Build model
 """
-function build_model_uc(instance; solver=Gurobi.Optimizer, PoolSearchMode=2, PoolSolutions=100)
+function build_model_uc(
+    instance;
+    solver = Gurobi.Optimizer,
+    PoolSearchMode = 2,
+    PoolSolutions = 100,
+)
     # Construct model (using state-of-the-art defaults)
-    model = UnitCommitment.build_model(
-        instance = instance,
-        optimizer = solver,
-    )
+    model = UnitCommitment.build_model(instance = instance, optimizer = solver)
 
     # Set solver attributes
     if !isnothing(PoolSearchMode)
@@ -54,7 +57,7 @@ end
 
 function tuple_2_name(smb)
     str = string(smb[1])
-    for i in 2:length(smb)
+    for i = 2:length(smb)
         str = str * "_" * string(smb[i])
     end
     return str
@@ -62,36 +65,30 @@ end
 
 function bin_variables_retriever(model)
     bin_vars = vcat(
-        collect(values(model[:is_on])), 
-        collect(values(model[:startup])), 
+        collect(values(model[:is_on])),
+        collect(values(model[:startup])),
         collect(values(model[:switch_on])),
-        collect(values(model[:switch_off]))
+        collect(values(model[:switch_off])),
     )
 
     bin_vars_names = vcat(
-        "is_on_" .* tuple_2_name.(collect(keys(model[:is_on]))), 
-        "startup_" .* tuple_2_name.(collect(keys(model[:startup]))), 
+        "is_on_" .* tuple_2_name.(collect(keys(model[:is_on]))),
+        "startup_" .* tuple_2_name.(collect(keys(model[:startup]))),
         "switch_on_" .* tuple_2_name.(collect(keys(model[:switch_on]))),
-        "switch_off_" .* tuple_2_name.(collect(keys(model[:switch_off])))
+        "switch_off_" .* tuple_2_name.(collect(keys(model[:switch_off]))),
     )
     return bin_vars, bin_vars_names
 end
 
 mutable struct StorageCallback <: Function
-    model
-    bin_vars
-    my_storage_vars
-    my_storage_obj
-    is_relaxed
+    model::Any
+    bin_vars::Any
+    my_storage_vars::Any
+    my_storage_obj::Any
+    is_relaxed::Any
 end
-function StorageCallback(model, bin_vars; maxiter=999999)
-    return StorageCallback(
-        model,
-        bin_vars,
-        [],
-        [],
-        [],
-    )
+function StorageCallback(model, bin_vars; maxiter = 999999)
+    return StorageCallback(model, bin_vars, [], [], [])
 end
 
 function (callback::StorageCallback)(cb_data, cb_where::Cint)
@@ -99,7 +96,10 @@ function (callback::StorageCallback)(cb_data, cb_where::Cint)
     if cb_where == GRB_CB_MIPNODE
         # prep
         obj_terms = objective_function(callback.model).terms
-        obj_terms_gurobi = [obj_terms[var] for var in all_variables(callback.model) if haskey(obj_terms, var)]
+        obj_terms_gurobi = [
+            obj_terms[var] for
+            var in all_variables(callback.model) if haskey(obj_terms, var)
+        ]
         num_all_var = num_variables(callback.model)
         # save
         resultobj = Ref{Cint}()
@@ -107,9 +107,19 @@ function (callback::StorageCallback)(cb_data, cb_where::Cint)
         if resultobj[] != GRB_OPTIMAL
             return  # Solution is something other than optimal.
         end
-        gurobi_indexes_all = [Gurobi.column(backend(callback.model).optimizer.model, callback.model.moi_backend.model_to_optimizer_map[var.index]) for var in all_variables(callback.model) if haskey(obj_terms, var)]
-        gurobi_indexes_bin = [Gurobi.column(backend(callback.model).optimizer.model, callback.model.moi_backend.model_to_optimizer_map[callback.bin_vars[i].index]) for i in 1:length(callback.bin_vars)]
-        
+        gurobi_indexes_all = [
+            Gurobi.column(
+                backend(callback.model).optimizer.model,
+                callback.model.moi_backend.model_to_optimizer_map[var.index],
+            ) for var in all_variables(callback.model) if haskey(obj_terms, var)
+        ]
+        gurobi_indexes_bin = [
+            Gurobi.column(
+                backend(callback.model).optimizer.model,
+                callback.model.moi_backend.model_to_optimizer_map[callback.bin_vars[i].index],
+            ) for i = 1:length(callback.bin_vars)
+        ]
+
         resultP = Vector{Cdouble}(undef, num_all_var)
         GRBcbget(cb_data, cb_where, GRB_CB_MIPNODE_REL, resultP)
         push!(callback.my_storage_vars, resultP[gurobi_indexes_bin])
@@ -138,7 +148,14 @@ end
 """
     Build branch and bound dataset
 """
-function uc_bnb_dataset(instance, save_file; model=build_model_uc(instance), data_dir=pwd(), batch_id = uuid1(), filetype=ArrowFile)
+function uc_bnb_dataset(
+    instance,
+    save_file;
+    model = build_model_uc(instance),
+    data_dir = pwd(),
+    batch_id = uuid1(),
+    filetype = ArrowFile,
+)
     ##############
     # Solve and store solutions
     ##############
@@ -172,22 +189,23 @@ function uc_bnb_dataset(instance, save_file; model=build_model_uc(instance), dat
     my_storage_obj = my_callback_function.my_storage_obj
 
     # Data
-    X = hcat(my_storage_vars...)'[:,:]
-    y = convert.(Float64, my_storage_obj[:,:])
+    X = hcat(my_storage_vars...)'[:, :]
+    y = convert.(Float64, my_storage_obj[:, :])
 
     # Save solutions
     # Input
-    instances_ids = [uuid1() for i in 1:length(my_storage_vars)]
+    instances_ids = [uuid1() for i = 1:length(my_storage_vars)]
     df_in = DataFrame(X, Symbol.(bin_vars_names))
     df_in.id = instances_ids
     # Save Loads
     for bus in instance.buses
-        for t in 1:instance.time
-            df_in[!, Symbol("load_" * string(bus.name) * "_" * string(t))] = fill(bus.load[t], length(instances_ids))
+        for t = 1:instance.time
+            df_in[!, Symbol("load_" * string(bus.name) * "_" * string(t))] =
+                fill(bus.load[t], length(instances_ids))
         end
     end
     # Output
-    df_out = DataFrame(Dict(:objective => y[:,1]))
+    df_out = DataFrame(Dict(:objective => y[:, 1]))
     df_out.id = instances_ids
     df_out.time = fill(solve_time(model), length(instances_ids))
     df_out.status = fill("LOCALLY_SOLVED", length(instances_ids))
@@ -202,11 +220,23 @@ function uc_bnb_dataset(instance, save_file; model=build_model_uc(instance), dat
 
     # Save
     if filetype === ArrowFile
-        Arrow.write(joinpath(data_dir, save_file * "_input_" * string(batch_id) * ".arrow"), df_in)
-        Arrow.write(joinpath(data_dir, save_file * "_output_" * string(batch_id) * ".arrow"), df_out)
+        Arrow.write(
+            joinpath(data_dir, save_file * "_input_" * string(batch_id) * ".arrow"),
+            df_in,
+        )
+        Arrow.write(
+            joinpath(data_dir, save_file * "_output_" * string(batch_id) * ".arrow"),
+            df_out,
+        )
     else
-        CSV.write(joinpath(data_dir, save_file * "_input_" * string(batch_id) * ".csv"), df_in)
-        CSV.write(joinpath(data_dir, save_file * "_output_" * string(batch_id) * ".csv"), df_out)
+        CSV.write(
+            joinpath(data_dir, save_file * "_input_" * string(batch_id) * ".csv"),
+            df_in,
+        )
+        CSV.write(
+            joinpath(data_dir, save_file * "_output_" * string(batch_id) * ".csv"),
+            df_out,
+        )
     end
 
     @info "Saved dataset to $(data_dir)" batch_id length(instances_ids) length(is_relaxed) optimal_obj
@@ -217,14 +247,25 @@ end
 """
     Enhance dataset
 """
-function uc_random_dataset!(instance, save_file; model=build_model_uc(instance), delete_objective=false, inner_solver=() -> POI.Optimizer(Gurobi.Optimizer()), data_dir=pwd(), filetype=ArrowFile, num_s = 1000, non_zero_units = 0.15, batch_id = uuid1())
+function uc_random_dataset!(
+    instance,
+    save_file;
+    model = build_model_uc(instance),
+    delete_objective = false,
+    inner_solver = () -> POI.Optimizer(Gurobi.Optimizer()),
+    data_dir = pwd(),
+    filetype = ArrowFile,
+    num_s = 1000,
+    non_zero_units = 0.15,
+    batch_id = uuid1(),
+)
     MOI.set(model, Gurobi.CallbackFunction(), nothing)
     bin_vars, bin_vars_names = bin_variables_retriever(model)
     # Remove binary constraints
     upper_model, inner_2_upper_map, cons_mapping = copy_binary_model(model)
 
     # delete binary constraints from inner model
-    delete_binary_terms!(model; delete_objective=delete_objective)
+    delete_binary_terms!(model; delete_objective = delete_objective)
     # add deficit constraints
     add_deficit_constraints!(model)
     # link binary variables from upper to inner model
@@ -236,10 +277,15 @@ function uc_random_dataset!(instance, save_file; model=build_model_uc(instance),
     # set solver
     set_optimizer(model, inner_solver)
     # Parameter values
-    u_values = abs.(Matrix(hcat([sprandn(length(u_inner), 1, non_zero_units) for i in 1:num_s]...)'))
+    u_values =
+        abs.(
+            Matrix(hcat([sprandn(length(u_inner), 1, non_zero_units) for i = 1:num_s]...)')
+        )
     u_values = min.(u_values, 1.0)
-    units_on = sum(u_values, dims=1)
-    @info "Number of units on: " mean(units_on) std(units_on) minimum(units_on) maximum(units_on)
+    units_on = sum(u_values, dims = 1)
+    @info "Number of units on: " mean(units_on) std(units_on) minimum(units_on) maximum(
+        units_on,
+    )
     parameter_values = Dict(u_inner .=> Array.(eachcol(u_values)))
     # The iterator
     problem_iterator = ProblemIterator(parameter_values)
@@ -256,8 +302,9 @@ function uc_random_dataset!(instance, save_file; model=build_model_uc(instance),
     rm(joinpath(data_dir, input_file))
     # add loads
     for bus in instance.buses
-        for t in 1:instance.time
-            df_in[!, Symbol("load_" * string(bus.name) * "_" * string(t))] = fill(bus.load[t], length(df_in.id))
+        for t = 1:instance.time
+            df_in[!, Symbol("load_" * string(bus.name) * "_" * string(t))] =
+                fill(bus.load[t], length(df_in.id))
         end
     end
     # save
@@ -268,9 +315,9 @@ function uc_random_dataset!(instance, save_file; model=build_model_uc(instance),
     end
     # CSV recorder to save the optimal primal and dual decision values
     output_file = save_file * "_output_" * string(batch_id)
-    recorder = Recorder{filetype}(joinpath(data_dir, output_file); model=model)
+    recorder = Recorder{filetype}(joinpath(data_dir, output_file); model = model)
     output_file = output_file * "." * string(filetype)
 
     # Finally solve all problems described by the iterator
     solve_batch(problem_iterator, recorder)
-end
\ No newline at end of file
+end
diff --git a/examples/unitcommitment/compress_arrow.jl b/examples/unitcommitment/compress_arrow.jl
index 3bde2ca..5c82c2e 100644
--- a/examples/unitcommitment/compress_arrow.jl
+++ b/examples/unitcommitment/compress_arrow.jl
@@ -5,20 +5,17 @@ using DataFrames
 case_name = "case300"
 date = "2017-01-01"
 horizon = 2
-case_file_path = joinpath(
-    dirname(@__FILE__), "data", case_name, date, "h" * string(horizon)
-)
+case_file_path =
+    joinpath(dirname(@__FILE__), "data", case_name, date, "h" * string(horizon))
 
 # Load input and output data tables
 iter_files = readdir(case_file_path)
 
-iter_files = [
-    file for file in iter_files if occursin(case_name, file) && occursin("arrow", file)
-]
+iter_files =
+    [file for file in iter_files if occursin(case_name, file) && occursin("arrow", file)]
 
-file_ins = [
-    joinpath(case_file_path, file) for file in iter_files if occursin("input", file)
-]
+file_ins =
+    [joinpath(case_file_path, file) for file in iter_files if occursin("input", file)]
 
 batch_ids = [split(split(file, "_")[end], ".")[1] for file in file_ins]
 
@@ -29,22 +26,28 @@ file_name = split(split(file_ins[1], "_input")[1], "/")[end]
 # move input files to input folder
 for file in iter_files
     if occursin("input", file)
-        mv(joinpath(case_file_path, file), joinpath(case_file_path, "input", file), force=true)
+        mv(
+            joinpath(case_file_path, file),
+            joinpath(case_file_path, "input", file),
+            force = true,
+        )
     end
 end
 
 # compress output files per batch id
 for batch_id in batch_ids
     file_outs_names = [
-        file
-        for file in iter_files
-        if occursin("output", file) && occursin(batch_id, file)
+        file for file in iter_files if occursin("output", file) && occursin(batch_id, file)
     ]
 
-    file_outs = [ joinpath(case_file_path, file) for file in file_outs_names ]
-    
+    file_outs = [joinpath(case_file_path, file) for file in file_outs_names]
+
     if length(file_outs_names) == 1
-        mv(file_outs[1], joinpath(case_file_path, "output", file_outs_names[1]), force=true)
+        mv(
+            file_outs[1],
+            joinpath(case_file_path, "output", file_outs_names[1]),
+            force = true,
+        )
         continue
     end
 
diff --git a/examples/unitcommitment/uc_dataset_generation.jl b/examples/unitcommitment/uc_dataset_generation.jl
index fca6a39..05bd72d 100644
--- a/examples/unitcommitment/uc_dataset_generation.jl
+++ b/examples/unitcommitment/uc_dataset_generation.jl
@@ -40,9 +40,7 @@ mkpath(joinpath(data_dir, "output"))
 ##############
 # Load Instance
 ##############
-instance = UnitCommitment.read_benchmark(
-    joinpath("matpower", case_name, date),
-)
+instance = UnitCommitment.read_benchmark(joinpath("matpower", case_name, date))
 instance.time = horizon
 
 ##############
@@ -51,35 +49,47 @@ instance.time = horizon
 
 if solve_nominal
     model = build_model_uc(instance)
-    uc_bnb_dataset(instance, save_file; data_dir=data_dir, model=model)
+    uc_bnb_dataset(instance, save_file; data_dir = data_dir, model = model)
     if num_random > 0
-        uc_random_dataset!(instance, save_file; data_dir=data_dir, model=model, num_s=num_random)
+        uc_random_dataset!(
+            instance,
+            save_file;
+            data_dir = data_dir,
+            model = model,
+            num_s = num_random,
+        )
     end
 end
 
 # save nominal loads in a dictionary
 nominal_loads = Dict()
-for i in 1:length(instance.buses)
+for i = 1:length(instance.buses)
     bus = instance.buses[i]
     nominal_loads[i] = bus.load[1:horizon]
 end
 
-@sync @distributed for i in 1:num_batches
+@sync @distributed for i = 1:num_batches
     rng = MersenneTwister(round(Int, i * time()))
     instance_ = deepcopy(instance)
     uc_load_disturbances!(rng, instance_, nominal_loads)
     # perturbed loads
     perturbed_loads_sum = 0.0
-    for i in 1:length(instance_.buses)
+    for i = 1:length(instance_.buses)
         bus = instance_.buses[i]
         perturbed_loads_sum += sum(bus.load)
     end
     @info "Solving batch $i" rng perturbed_loads_sum
     model = build_model_uc(instance_)
-    uc_bnb_dataset(instance_, save_file; data_dir=data_dir, model=model)
+    uc_bnb_dataset(instance_, save_file; data_dir = data_dir, model = model)
     if num_random > 0
-        uc_random_dataset!(instance_, save_file; data_dir=data_dir, model=model, num_s=num_random)
+        uc_random_dataset!(
+            instance_,
+            save_file;
+            data_dir = data_dir,
+            model = model,
+            num_s = num_random,
+        )
     end
 end
 
-include(joinpath(dirname(@__FILE__), "compress_arrow.jl"))
\ No newline at end of file
+include(joinpath(dirname(@__FILE__), "compress_arrow.jl"))
diff --git a/examples/unitcommitment/uc_nn_training.jl b/examples/unitcommitment/uc_nn_training.jl
index e44e489..ca56ba9 100644
--- a/examples/unitcommitment/uc_nn_training.jl
+++ b/examples/unitcommitment/uc_nn_training.jl
@@ -6,7 +6,9 @@
 # Load Functions
 ##############
 
-import Pkg; Pkg.activate(dirname(dirname(@__DIR__))); Pkg.instantiate()
+import Pkg;
+Pkg.activate(dirname(dirname(@__DIR__)));
+Pkg.instantiate();
 
 using LearningToOptimize
 using MLJFlux
@@ -29,7 +31,7 @@ data_dir = joinpath(dirname(@__FILE__), "data")
 ##############
 # Parameters
 ##############
-filetype=ArrowFile
+filetype = ArrowFile
 case_name = ARGS[1] # case_name = "case300"
 date = ARGS[2] # date="2017-01-01"
 horizon = parse(Int, ARGS[3]) # horizon=2
@@ -41,11 +43,13 @@ data_dir = joinpath(data_dir, case_name, date, "h" * string(horizon))
 ##############
 
 # Read input and output data
-iter_input = readdir(joinpath(data_dir, "input"), join=true)
-iter_output = readdir(joinpath(data_dir, "output"), join=true)
+iter_input = readdir(joinpath(data_dir, "input"), join = true)
+iter_output = readdir(joinpath(data_dir, "output"), join = true)
 # filter for only arrow files of this case
-iter_input = [file for file in iter_input if occursin(case_name, file) && occursin("arrow", file)]
-iter_output = [file for file in iter_output if occursin(case_name, file) && occursin("arrow", file)]
+iter_input =
+    [file for file in iter_input if occursin(case_name, file) && occursin("arrow", file)]
+iter_output =
+    [file for file in iter_output if occursin(case_name, file) && occursin("arrow", file)]
 
 # Load input and output data tables
 input_tables = Array{DataFrame}(undef, length(iter_input))
@@ -62,7 +66,7 @@ input_table = vcat(input_tables...)
 output_table = vcat(output_tables...)
 
 # Separate input and output variables & ignore id time status primal_status dual_status
-train_table = innerjoin(input_table, output_table[!, [:id, :objective]]; on=:id)
+train_table = innerjoin(input_table, output_table[!, [:id, :objective]]; on = :id)
 input_features = names(train_table[!, Not([:id, :objective])])
 X = Float32.(Matrix(train_table[!, input_features]))
 y = Float32.(Matrix(train_table[!, [:objective]]))
@@ -79,22 +83,27 @@ lg = WandbLogger(
         "learning_rate" => 0.01,
         "rng" => 123,
         # "lambda" => 0.00,
-    )
+    ),
 )
 
-optimiser=ConvexRule(
-    Flux.Optimise.Adam(get_config(lg, "learning_rate"), (0.9, 0.999), 1.0e-8, IdDict{Any,Any}())
+optimiser = ConvexRule(
+    Flux.Optimise.Adam(
+        get_config(lg, "learning_rate"),
+        (0.9, 0.999),
+        1.0e-8,
+        IdDict{Any,Any}(),
+    ),
 )
 
 nn = MultitargetNeuralNetworkRegressor(;
-    builder=FullyConnectedBuilder(layers),
-    rng=get_config(lg, "rng"),
-    epochs=5000,
-    optimiser=optimiser,
-    acceleration=CUDALibs(),
-    batch_size=get_config(lg, "batch_size"),
+    builder = FullyConnectedBuilder(layers),
+    rng = get_config(lg, "rng"),
+    epochs = 5000,
+    optimiser = optimiser,
+    acceleration = CUDALibs(),
+    batch_size = get_config(lg, "batch_size"),
     # lambda=get_config(lg, "lambda"),
-    loss=relative_rmse,
+    loss = relative_rmse,
 )
 
 # Constrols
@@ -103,18 +112,19 @@ model_dir = joinpath(dirname(@__FILE__), "models")
 mkpath(model_dir)
 
 save_control =
-    MLJIteration.skip(Save(joinpath(model_dir, save_file * ".jls")), predicate=3)
+    MLJIteration.skip(Save(joinpath(model_dir, save_file * ".jls")), predicate = 3)
 
-controls=[Step(2),
+controls = [
+    Step(2),
     NumberSinceBest(6),
     # PQ(; alpha=0.9, k=30),
-    GL(; alpha=4.0),
+    GL(; alpha = 4.0),
     InvalidValue(),
-    TimeLimit(; t=1),
+    TimeLimit(; t = 1),
     save_control,
     WithLossDo(update_loss),
     WithReportDo(update_training_loss),
-    WithIterationsDo(update_epochs)
+    WithIterationsDo(update_epochs),
 ]
 
 # WIP
@@ -126,17 +136,17 @@ controls=[Step(2),
 # layer = mach.fitresult.fitresult[1]
 # gradient( ( _x ) -> sum(layer( _x )), X')
 
-iterated_pipe =
-    IteratedModel(model=nn,
-        controls=controls,
-        resampling=Holdout(fraction_train=0.7),
-        measure = relative_mae,
+iterated_pipe = IteratedModel(
+    model = nn,
+    controls = controls,
+    resampling = Holdout(fraction_train = 0.7),
+    measure = relative_mae,
 )
 
 # Fit model
 # clear()
 mach = machine(iterated_pipe, X, y)
-fit!(mach; verbosity=2)
+fit!(mach; verbosity = 2)
 
 # Finish the run
 close(lg)
@@ -152,4 +162,9 @@ fitted_model = mach.fitresult.fitresult[1]
 
 model_state = Flux.state(fitted_model)
 
-jldsave(joinpath(model_dir, save_file * ".jld2"); model_state=model_state, layers=layers, input_features=input_features)
+jldsave(
+    joinpath(model_dir, save_file * ".jld2");
+    model_state = model_state,
+    layers = layers,
+    input_features = input_features,
+)
diff --git a/examples/unitcommitment/unitcommitment.jl b/examples/unitcommitment/unitcommitment.jl
index b6d4536..899fe7e 100644
--- a/examples/unitcommitment/unitcommitment.jl
+++ b/examples/unitcommitment/unitcommitment.jl
@@ -30,7 +30,7 @@ horizon = 2
 save_file = case_name * "_" * replace(date, "-" => "_") * "_h" * string(horizon)
 model_dir = joinpath(pwd(), "examples", "unitcommitment", "models")
 data_file_path = joinpath(pwd(), "examples", "unitcommitment", "data")
-case_file_path = joinpath(data_file_path, case_name, date, "h"*string(horizon))
+case_file_path = joinpath(data_file_path, case_name, date, "h" * string(horizon))
 
 upper_solver = Gurobi.Optimizer
 
@@ -51,33 +51,31 @@ layers = model_save["layers"]
 # Load Data
 ##############
 
-file_in = readdir(joinpath(case_file_path, "input"), join=true)[1]
+file_in = readdir(joinpath(case_file_path, "input"), join = true)[1]
 batch_id = split(split(split(file_in, "/")[end], "_input_")[2], ".")[1]
-file_outs = readdir(joinpath(case_file_path, "output"), join=true)
+file_outs = readdir(joinpath(case_file_path, "output"), join = true)
 file_out = [file for file in file_outs if occursin(batch_id, file)][1]
 
 input_table = Arrow.Table(file_in) |> DataFrame
 output_table = Arrow.Table(file_out) |> DataFrame
 
-train_table = innerjoin(input_table, output_table[!, [:id, :objective]]; on=:id)
+train_table = innerjoin(input_table, output_table[!, [:id, :objective]]; on = :id)
 X = Float32.(Matrix(train_table[!, Symbol.(input_features)]))
 y = Float32.(Matrix(train_table[!, [:objective]]))
 
-true_ob_value = output_table[output_table.status .== "OPTIMAL", :objective][1]
+true_ob_value = output_table[output_table.status.=="OPTIMAL", :objective][1]
 
 ##############
 # Load Instance
 ##############
-instance = UnitCommitment.read_benchmark(
-    joinpath("matpower", case_name, date),
-)
+instance = UnitCommitment.read_benchmark(joinpath("matpower", case_name, date))
 instance.time = horizon
 
 # impose load
-for i in 1:length(instance.buses)
+for i = 1:length(instance.buses)
     bus = instance.buses[i]
     bus_name = bus.name
-    for h in 1:horizon
+    for h = 1:horizon
         bus.load[h] = input_table[1, Symbol("load_" * bus_name * "_" * string(h))]
     end
 end
@@ -111,7 +109,7 @@ end
 
 # add nn to model
 bin_vars_upper = Array{Any}(undef, length(input_features))
-for i in 1:length(input_features)
+for i = 1:length(input_features)
     feature = input_features[i]
     if occursin("load", feature)
         bin_vars_upper[i] = input_table[1, Symbol(feature)]
@@ -148,4 +146,4 @@ for var in bin_vars
 end
 
 JuMP.optimize!(model)
-abs(objective_value(model) - true_ob_value) / true_ob_value
\ No newline at end of file
+abs(objective_value(model) - true_ob_value) / true_ob_value
diff --git a/examples/unitcommitment/visualize.jl b/examples/unitcommitment/visualize.jl
index 8984238..1212545 100644
--- a/examples/unitcommitment/visualize.jl
+++ b/examples/unitcommitment/visualize.jl
@@ -4,29 +4,30 @@ using DataFrames
 using Statistics
 using LinearAlgebra
 
-cossim(x,y) = dot(x,y) / (norm(x)*norm(y))
+cossim(x, y) = dot(x, y) / (norm(x) * norm(y))
 
 # Data Parameters
 case_name = "case300"
 date = "2017-01-01"
 horizon = "2"
 path_dataset = joinpath(dirname(@__FILE__), "data")
-case_file_path = joinpath(path_dataset, case_name, date, "h"*horizon)
+case_file_path = joinpath(path_dataset, case_name, date, "h" * horizon)
 
 # Load input and output data tables
-file_ins = readdir(joinpath(case_file_path, "input"), join=true)
+file_ins = readdir(joinpath(case_file_path, "input"), join = true)
 batch_ids = [split(split(file, "_")[end], ".")[1] for file in file_ins]
-file_outs = readdir(joinpath(case_file_path, "output"), join=true)
-file_outs = [file_outs[findfirst(x -> occursin(batch_id, x), file_outs)] for batch_id in batch_ids]
+file_outs = readdir(joinpath(case_file_path, "output"), join = true)
+file_outs =
+    [file_outs[findfirst(x -> occursin(batch_id, x), file_outs)] for batch_id in batch_ids]
 
 # Load input and output data tables
 input_tables = Array{DataFrame}(undef, length(file_ins))
 output_tables = Array{DataFrame}(undef, length(file_outs))
-for i in 1:length(file_ins)
+for i = 1:length(file_ins)
     file = file_ins[i]
     input_tables[i] = Arrow.Table(file) |> DataFrame
 end
-for i in 1:length(file_outs)
+for i = 1:length(file_outs)
     file = file_outs[i]
     output_tables[i] = Arrow.Table(file) |> DataFrame
     # if all the status are OPTIMAL, make them INFEASIBLE
@@ -40,73 +41,107 @@ load_columns = [col for col in names(input_tables[1]) if occursin("load", col)]
 nominal_loads = Vector(input_tables[1][1, load_columns])
 
 # Load divergence
-cos_sim = [acos(cossim(nominal_loads, Vector(input_table[1, load_columns]))) for input_table in input_tables]
-norm_sim = [norm(nominal_loads) - norm(Vector(input_table[1, load_columns])) for input_table in input_tables] ./ norm(nominal_loads) .+ 1
+cos_sim = [
+    acos(cossim(nominal_loads, Vector(input_table[1, load_columns]))) for
+    input_table in input_tables
+]
+norm_sim =
+    [
+        norm(nominal_loads) - norm(Vector(input_table[1, load_columns])) for
+        input_table in input_tables
+    ] ./ norm(nominal_loads) .+ 1
 
 # Polar Plot divergence
 plotly()
-plt = plot(cos_sim, norm_sim, title="Load Similarity", proj = :polar, m = 2)
+plt = plot(cos_sim, norm_sim, title = "Load Similarity", proj = :polar, m = 2)
 
 # concatenate all the input and output tables
 input_table = vcat(input_tables...)
 output_table = vcat(output_tables...)
 
 # sum each row of the input table (excluding the id) and store it in a new column
-commit_columns = [col for col in names(input_table) if !occursin("load", col) && !occursin("id", col)]
+commit_columns =
+    [col for col in names(input_table) if !occursin("load", col) && !occursin("id", col)]
 input_table.total_commitment = sum(eachcol(input_table[!, commit_columns]))
 input_table.total_load = sum(eachcol(input_table[!, load_columns]))
 
 # join input and output tables keep on id and keep only total_commitment and objective
-table_train = innerjoin(input_table, output_table[!, [:id, :objective, :status]]; on=:id)[!, [:id, :total_commitment, :objective, :status, :total_load]]
+table_train = innerjoin(input_table, output_table[!, [:id, :objective, :status]]; on = :id)[
+    !,
+    [:id, :total_commitment, :objective, :status, :total_load],
+]
 
 # separate infeasible from feasible
-table_train_optimal = table_train[table_train.status .== "OPTIMAL", :]
-table_train_infeasible = table_train[table_train.status .== "INFEASIBLE", :]
-table_train_localopt = table_train[table_train.status .== "LOCALLY_SOLVED", :]
+table_train_optimal = table_train[table_train.status.=="OPTIMAL", :]
+table_train_infeasible = table_train[table_train.status.=="INFEASIBLE", :]
+table_train_localopt = table_train[table_train.status.=="LOCALLY_SOLVED", :]
 
 # plot objective vs total commitment
 # now with the y axis on the log scale
 plotly()
 plt = scatter(
-    table_train_optimal.total_commitment, 
-    table_train_optimal.objective, 
-    label="Optimal", xlabel="Total Commitment", ylabel="Objective", 
-    title="", color=:red, yscale=:log10, legend=:outertopright,
+    table_train_optimal.total_commitment,
+    table_train_optimal.objective,
+    label = "Optimal",
+    xlabel = "Total Commitment",
+    ylabel = "Objective",
+    title = "",
+    color = :red,
+    yscale = :log10,
+    legend = :outertopright,
     # marker size
-    ms=10,
+    ms = 10,
 );
-scatter!(plt, 
-    table_train_localopt.total_commitment, 
-    table_train_localopt.objective, 
-    label="Local Optimum", color=:blue,
-    marker=:o, alpha=0.5
+scatter!(
+    plt,
+    table_train_localopt.total_commitment,
+    table_train_localopt.objective,
+    label = "Local Optimum",
+    color = :blue,
+    marker = :o,
+    alpha = 0.5,
 );
-scatter!(plt, 
-    table_train_infeasible.total_commitment, 
-    table_train_infeasible.objective, 
-    label="Infeasible", color=:yellow,
-    marker=:square, alpha=0.01, ms=2
+scatter!(
+    plt,
+    table_train_infeasible.total_commitment,
+    table_train_infeasible.objective,
+    label = "Infeasible",
+    color = :yellow,
+    marker = :square,
+    alpha = 0.01,
+    ms = 2,
 )
 
 # plot objective vs total load
 # now with the y axis on the log scale
 plt2 = scatter(
-    table_train_optimal.total_load, 
-    table_train_optimal.objective, 
-    label="Optimal", xlabel="Total Load", ylabel="Objective", 
-    title="", color=:red, yscale=:log10, legend=:outertopright,
+    table_train_optimal.total_load,
+    table_train_optimal.objective,
+    label = "Optimal",
+    xlabel = "Total Load",
+    ylabel = "Objective",
+    title = "",
+    color = :red,
+    yscale = :log10,
+    legend = :outertopright,
     # marker size
-    ms=10,
+    ms = 10,
 );
-scatter!(plt2, 
-    table_train_localopt.total_load, 
-    table_train_localopt.objective, 
-    label="Local Optimum", color=:blue,
-    marker=:o, alpha=0.5
+scatter!(
+    plt2,
+    table_train_localopt.total_load,
+    table_train_localopt.objective,
+    label = "Local Optimum",
+    color = :blue,
+    marker = :o,
+    alpha = 0.5,
 );
-scatter!(plt2, 
-    table_train_infeasible.total_load, 
-    table_train_infeasible.objective, 
-    label="Infeasible", color=:yellow,
-    marker=:square, alpha=0.5
+scatter!(
+    plt2,
+    table_train_infeasible.total_load,
+    table_train_infeasible.objective,
+    label = "Infeasible",
+    color = :yellow,
+    marker = :square,
+    alpha = 0.5,
 )
diff --git a/examples/worst_case_active_learning/test_worst_case.jl b/examples/worst_case_active_learning/test_worst_case.jl
index f1d7544..2ebd43c 100644
--- a/examples/worst_case_active_learning/test_worst_case.jl
+++ b/examples/worst_case_active_learning/test_worst_case.jl
@@ -5,7 +5,7 @@ using NonconvexNLopt
 
 Test dataset generation using the worst case problem iterator for different filetypes.
 """
-function test_worst_case_problem_iterator(path::AbstractString, num_p=10)
+function test_worst_case_problem_iterator(path::AbstractString, num_p = 10)
     @testset "Worst Case Dual Generation Type: $filetype" for filetype in
                                                               [CSVFile, ArrowFile]
         # The problem to iterate over
@@ -13,7 +13,7 @@ function test_worst_case_problem_iterator(path::AbstractString, num_p=10)
             return () -> Ipopt.Optimizer()
         end
         parameter_factory = (model) -> [@variable(model, _p)]
-        function primal_builder!(model, parameters; recorder=nothing)
+        function primal_builder!(model, parameters; recorder = nothing)
             @variable(model, x)
             @constraint(model, cons, x + parameters[1] >= 3)
             @objective(model, Min, 2x)
@@ -30,7 +30,7 @@ function test_worst_case_problem_iterator(path::AbstractString, num_p=10)
         end
 
         problem_iterator = WorstCaseProblemIterator(
-            [uuid1() for _ in 1:num_p],
+            [uuid1() for _ = 1:num_p],
             parameter_factory,
             primal_builder!,
             set_iterator!,
@@ -50,7 +50,10 @@ function test_worst_case_problem_iterator(path::AbstractString, num_p=10)
 
         # The recorder
         recorder = Recorder{filetype}(
-            file_output; filename_input=file_input, primal_variables=[model[:x]], dual_variables=[]
+            file_output;
+            filename_input = file_input,
+            primal_variables = [model[:x]],
+            dual_variables = [],
         )
 
         # Solve all problems and record solutions
@@ -96,7 +99,7 @@ function test_worst_case_problem_iterator(path::AbstractString, num_p=10)
 
     @testset "Worst Case NonConvex Generation Type: $filetype" for filetype in
                                                                    [CSVFile, ArrowFile]
-        function _primal_builder!(; recorder=nothing)
+        function _primal_builder!(; recorder = nothing)
             model = JuMP.Model(() -> POI.Optimizer(HiGHS.Optimizer()))
             parameters = @variable(model, _p in MOI.Parameter(1.0))
             @variable(model, x)
@@ -120,12 +123,12 @@ function test_worst_case_problem_iterator(path::AbstractString, num_p=10)
         end
 
         problem_iterator = WorstCaseProblemIterator(
-            [uuid1() for _ in 1:num_p], # will be ignored
+            [uuid1() for _ = 1:num_p], # will be ignored
             () -> nothing, # will be ignored
             _primal_builder!,
             _set_iterator!,
             NLoptAlg(:LN_BOBYQA);
-            options=NLoptOptions(; maxeval=10),
+            options = NLoptOptions(; maxeval = 10),
         )
 
         # Test Build Primal
@@ -138,7 +141,10 @@ function test_worst_case_problem_iterator(path::AbstractString, num_p=10)
 
         # The recorder
         recorder = Recorder{filetype}(
-            file_output; filename_input=file_input, primal_variables=[model[:x]], dual_variables=[]
+            file_output;
+            filename_input = file_input,
+            primal_variables = [model[:x]],
+            dual_variables = [],
         )
 
         # Solve all problems and record solutions
@@ -186,10 +192,6 @@ end
 ####### Run Tests
 path = mktempdir()
 test_worst_case_problem_iterator(path)
-file_in, file_out = test_generate_worst_case_dataset(
-    path, "pglib_opf_case5_pjm", 20
-)
-file_in, file_out = test_generate_worst_case_dataset_Nonconvex(
-    path, "pglib_opf_case5_pjm", 20
-)
-
+file_in, file_out = test_generate_worst_case_dataset(path, "pglib_opf_case5_pjm", 20)
+file_in, file_out =
+    test_generate_worst_case_dataset_Nonconvex(path, "pglib_opf_case5_pjm", 20)
diff --git a/examples/worst_case_active_learning/worst_case.jl b/examples/worst_case_active_learning/worst_case.jl
index cea1afe..3b673d5 100644
--- a/examples/worst_case_active_learning/worst_case.jl
+++ b/examples/worst_case_active_learning/worst_case.jl
@@ -38,10 +38,10 @@ struct WorstCaseProblemIterator{F} <: AbstractProblemIterator
         primal_builder!::Function,
         set_iterator!::Function,
         optimizer::F;
-        hook::Union{Nothing,Function}=nothing,
-        options::Any=nothing,
-        ext::Dict=Dict(),
-        early_stop::Function=(args...) -> false,
+        hook::Union{Nothing,Function} = nothing,
+        options::Any = nothing,
+        ext::Dict = Dict(),
+        early_stop::Function = (args...) -> false,
     ) where {F}
         return new{F}(
             ids,
@@ -69,7 +69,9 @@ solution if it passes the filter function.
 - `idx::Integer`: The index of the problem to solve.
 """
 function solve_and_record(
-    problem_iterator::WorstCaseProblemIterator, recorder::Recorder, idx::Integer
+    problem_iterator::WorstCaseProblemIterator,
+    recorder::Recorder,
+    idx::Integer,
 )
     # Build Primal
     model = JuMP.Model()
@@ -80,7 +82,7 @@ function solve_and_record(
     load_moi_idx = Vector{MOI.VariableIndex}(JuMP.index.(parameters))
 
     # Dualize the model
-    dual_st = Dualization.dualize(JuMP.backend(model); variable_parameters=load_moi_idx)
+    dual_st = Dualization.dualize(JuMP.backend(model); variable_parameters = load_moi_idx)
 
     dual_model = dual_st.dual_model
     primal_dual_map = dual_st.primal_dual_map
@@ -94,9 +96,8 @@ function solve_and_record(
     end
 
     # Get dual variables for the parameters
-    load_dual_idxs = [
-        map_moi_to_jump[primal_dual_map.primal_parameter[l]].value for l in load_moi_idx
-    ]
+    load_dual_idxs =
+        [map_moi_to_jump[primal_dual_map.primal_parameter[l]].value for l in load_moi_idx]
     load_var_dual = JuMP.all_variables(jump_dual_model)[load_dual_idxs]
 
     # Add constraints to the dual associated with the parameters
@@ -124,7 +125,7 @@ function solve_and_record(
     if recorder.filterfn(jump_dual_model)
         recorder.primal_variables = load_var_dual
         recorder.dual_variables = []
-        record(recorder, problem_iterator.ids[idx]; input=true)
+        record(recorder, problem_iterator.ids[idx]; input = true)
     else
         return 0, early_stop_bool
     end
@@ -134,7 +135,7 @@ function solve_and_record(
 
     # Create final primal model and solve
     model = JuMP.Model(problem_iterator.optimizer())
-    problem_iterator.primal_builder!(model, optimal_loads; recorder=recorder)
+    problem_iterator.primal_builder!(model, optimal_loads; recorder = recorder)
     JuMP.optimize!(model)
 
     termination_status = recorder.filterfn(model)
@@ -148,7 +149,7 @@ function solve_and_record(
     solution_primal_status != MOI.FEASIBLE_POINT && @warn("Primal solution not found")
     solution_dual_status != MOI.FEASIBLE_POINT && @warn("Dual solution not found")
 
-    if !isapprox(optimal_final_cost, optimal_dual_cost; rtol=1e-4)
+    if !isapprox(optimal_final_cost, optimal_dual_cost; rtol = 1e-4)
         rtol = abs(optimal_final_cost - optimal_dual_cost) / optimal_final_cost * 100
         @warn "Final cost is not equal to dual cost by $(rtol) %" optimal_final_cost optimal_dual_cost
     end
diff --git a/examples/worst_case_active_learning/worst_case_iter.jl b/examples/worst_case_active_learning/worst_case_iter.jl
index 45cca53..7053810 100644
--- a/examples/worst_case_active_learning/worst_case_iter.jl
+++ b/examples/worst_case_active_learning/worst_case_iter.jl
@@ -1,7 +1,7 @@
 using Nonconvex
 
 # Nonconvex needs a minimization objective function that only receives the decision vector.
-function primal_objective(parameter_values, parameters, filter_fn; penalty=1e8)
+function primal_objective(parameter_values, parameters, filter_fn; penalty = 1e8)
     model = owner_model(first(parameters))
     for (i, p) in enumerate(parameters)
         update_model!(model, p, parameter_values[i])
@@ -30,7 +30,7 @@ function save_input(parameters, _recorder, id)
     recorder = similar(_recorder)
     set_primal_variable!(recorder, parameters)
     set_dual_variable!(recorder, [])
-    record(recorder, id; input=true)
+    record(recorder, id; input = true)
     return nothing
 end
 
@@ -66,17 +66,15 @@ function solve_and_record(
     problem_iterator::WorstCaseProblemIterator{T},
     recorder::Recorder,
     idx::Integer;
-    maxiter=1000,
+    maxiter = 1000,
 ) where {T<:NonconvexCore.AbstractOptimizer}
     # Build Primal
-    model, parameters = problem_iterator.primal_builder!(; recorder=recorder)
-    (min_demands, max_demands, max_total_volume, starting_point) = problem_iterator.set_iterator!(
-        idx
-    )
+    model, parameters = problem_iterator.primal_builder!(; recorder = recorder)
+    (min_demands, max_demands, max_total_volume, starting_point) =
+        problem_iterator.set_iterator!(idx)
 
-    storage_objective_function = StorageCallbackObjective(
-        parameters, recorder.filterfn, recorder
-    )
+    storage_objective_function =
+        StorageCallbackObjective(parameters, recorder.filterfn, recorder)
 
     if haskey(problem_iterator.ext, :best_solution)
         starting_point = problem_iterator.ext[:best_solution]
@@ -84,7 +82,7 @@ function solve_and_record(
 
     # Build Nonconvex optimization model:
     model_non = Nonconvex.Model()
-    set_objective!(model_non, storage_objective_function; flags=[:expensive])
+    set_objective!(model_non, storage_objective_function; flags = [:expensive])
     addvar!(model_non, min_demands, max_demands)
     # add_ineq_constraint!(model_non, x -> sum(x .^ 2) - max_total_volume ^ 2)
 
@@ -94,7 +92,7 @@ function solve_and_record(
             model_non,
             problem_iterator.optimizer,
             starting_point;
-            options=problem_iterator.options,
+            options = problem_iterator.options,
         )
     else
         optimize(model_non, problem_iterator.optimizer, starting_point)
diff --git a/src/FullyConnected.jl b/src/FullyConnected.jl
index 67fbef6..4478b6b 100644
--- a/src/FullyConnected.jl
+++ b/src/FullyConnected.jl
@@ -27,25 +27,25 @@ function FullyConnected(
     input_size::Int,
     hidden_sizes::Vector{Int},
     output_size::Int;
-    init=Flux.glorot_uniform(Random.GLOBAL_RNG),
+    init = Flux.glorot_uniform(Random.GLOBAL_RNG),
 )::FullyConnected
     # Create layers
     layers = []
 
     # Create the pass through layers
-    pass_through = [Dense(input_size, 1; init=init) for _ in 2:(length(hidden_sizes) + 1)]
+    pass_through = [Dense(input_size, 1; init = init) for _ = 2:(length(hidden_sizes)+1)]
 
     # Create the first layer connected to the input size
-    push!(layers, Dense(input_size, hidden_sizes[1], relu; init=init))
+    push!(layers, Dense(input_size, hidden_sizes[1], relu; init = init))
 
     # Create connections between hidden layers
-    for i in 2:length(hidden_sizes)
+    for i = 2:length(hidden_sizes)
         # Create a new layer
-        push!(layers, Dense(hidden_sizes[i - 1] + 1, hidden_sizes[i], relu; init=init))
+        push!(layers, Dense(hidden_sizes[i-1] + 1, hidden_sizes[i], relu; init = init))
     end
 
     # Create the output layer connected to the last hidden layer
-    push!(layers, Dense(hidden_sizes[end] + 1, output_size; init=init))
+    push!(layers, Dense(hidden_sizes[end] + 1, output_size; init = init))
 
     return FullyConnected(PairwiseFusion(vcat, layers...), pass_through) |> gpu
 end
@@ -56,7 +56,7 @@ end
 
 function MLJFlux.build(builder::FullyConnectedBuilder, rng, n_in, n_out)
     init = Flux.glorot_uniform(rng)
-    return Chain(FullyConnected(n_in, builder.hidden_sizes, n_out; init=init)) |> gpu
+    return Chain(FullyConnected(n_in, builder.hidden_sizes, n_out; init = init)) |> gpu
 end
 
 # mutable struct ConvexRegressor <: MLJFlux.MLJFluxDeterministic
@@ -70,13 +70,13 @@ struct ConvexRule <: Optimisers.AbstractRule
     rule::Optimisers.AbstractRule
     tol::Real
 end
-function ConvexRule(rule::Optimisers.AbstractRule; tol=1e-6)
+function ConvexRule(rule::Optimisers.AbstractRule; tol = 1e-6)
     return ConvexRule(rule, tol)
 end
 
 Optimisers.init(o::ConvexRule, x::AbstractArray) = Optimisers.init(o.rule, x)
 
-function Optimisers.apply!(o::ConvexRule, mvel, x::AbstractArray{T}, dx) where T
+function Optimisers.apply!(o::ConvexRule, mvel, x::AbstractArray{T}, dx) where {T}
     return Optimisers.apply!(o.rule, mvel, x, dx)
 end
 
@@ -86,7 +86,7 @@ end
 Make a `PairwiseFusion` model convex by making sure all dense layers (a part from the first) have positive weights.
 This procedure only makes sense for single output fully connected models.
 """
-function make_convex!(chain::PairwiseFusion; tol=1e-6)
+function make_convex!(chain::PairwiseFusion; tol = 1e-6)
     for layer in chain.layers[2:end]
         layer.weight .= max.(layer.weight, tol)
     end
@@ -98,13 +98,13 @@ end
 Make a `FullyConnected` model convex by making sure all dense layers (a part from the first) have positive weights.
 This procedure only makes sense for single output fully connected models.
 """
-function make_convex!(model::FullyConnected; tol=1e-6)
-    return make_convex!(model.layers; tol=tol)
+function make_convex!(model::FullyConnected; tol = 1e-6)
+    return make_convex!(model.layers; tol = tol)
 end
 
-function make_convex!(model::Chain; tol=1e-6)
-    for i in 1:length(model.layers)
-        make_convex!(model.layers[i]; tol=tol)
+function make_convex!(model::Chain; tol = 1e-6)
+    for i = 1:length(model.layers)
+        make_convex!(model.layers[i]; tol = tol)
     end
 end
 
@@ -117,33 +117,32 @@ function MLJFlux.train(
     verbosity,
     X,
     y,
-    )
+)
 
     loss = model.loss
 
     # intitialize and start progress meter:
-    meter = MLJFlux.Progress(epochs + 1, dt=0, desc="Optimising neural net:",
-        barglyphs=MLJFlux.BarGlyphs("[=> ]"), barlen=25, color=:yellow)
+    meter = MLJFlux.Progress(
+        epochs + 1,
+        dt = 0,
+        desc = "Optimising neural net:",
+        barglyphs = MLJFlux.BarGlyphs("[=> ]"),
+        barlen = 25,
+        color = :yellow,
+    )
     verbosity != 1 || MLJFlux.next!(meter)
 
     # initiate history:
     n_batches = length(y)
 
-    losses = (loss(chain(X[i]), y[i]) for i in 1:n_batches)
-    history = [mean(losses),]
-
-    for i in 1:epochs
-        chain, optimiser_state, current_loss = MLJFlux.train_epoch(
-            model,
-            chain,
-            optimiser,
-            optimiser_state,
-            X,
-            y,
-        )
-        make_convex!(chain; tol=optimiser.tol)
-        verbosity < 2 ||
-            @info "Loss is $(round(current_loss; sigdigits=4))"
+    losses = (loss(chain(X[i]), y[i]) for i = 1:n_batches)
+    history = [mean(losses)]
+
+    for i = 1:epochs
+        chain, optimiser_state, current_loss =
+            MLJFlux.train_epoch(model, chain, optimiser, optimiser_state, X, y)
+        make_convex!(chain; tol = optimiser.tol)
+        verbosity < 2 || @info "Loss is $(round(current_loss; sigdigits=4))"
         verbosity != 1 || next!(meter)
         push!(history, current_loss)
     end
@@ -152,21 +151,19 @@ function MLJFlux.train(
 
 end
 
-function train!(model, loss, opt_state, X, Y; _batchsize=32, shuffle=true)
+function train!(model, loss, opt_state, X, Y; _batchsize = 32, shuffle = true)
     batchsize = min(size(X, 2), _batchsize)
     X = X |> gpu
     Y = Y |> gpu
-    data = Flux.DataLoader((X, Y), 
-        batchsize=batchsize, shuffle=shuffle
-    )
+    data = Flux.DataLoader((X, Y), batchsize = batchsize, shuffle = shuffle)
     for d in data
-		∇model, _ = gradient(model, d...) do m, x, y  # calculate the gradients
-			loss(m(x), y)
-		end;
-		# insert what ever code you want here that needs gradient
-		# E.g. logging with TensorBoardLogger.jl as histogram so you can see if it is becoming huge
-		opt_state, model = Optimisers.update(opt_state, model, ∇model)
-		# Here you might like to check validation set accuracy, and break out to do early stopping
-	end
+        ∇model, _ = gradient(model, d...) do m, x, y  # calculate the gradients
+            loss(m(x), y)
+        end
+        # insert what ever code you want here that needs gradient
+        # E.g. logging with TensorBoardLogger.jl as histogram so you can see if it is becoming huge
+        opt_state, model = Optimisers.update(opt_state, model, ∇model)
+        # Here you might like to check validation set accuracy, and break out to do early stopping
+    end
     return loss(model(X), Y)
 end
diff --git a/src/LearningToOptimize.jl b/src/LearningToOptimize.jl
index e76d6af..af24be4 100644
--- a/src/LearningToOptimize.jl
+++ b/src/LearningToOptimize.jl
@@ -7,12 +7,14 @@ using JuMP
 using UUIDs
 import ParametricOptInterface as POI
 import JuMP.MOI as MOI
+import JuMP.MOI.Utilities: distance_to_set
 import Base: string
 using Statistics
 using Distributions
 
 using Zygote
 
+using MLJ
 using MLJFlux
 using Flux
 using Flux: @functor
@@ -42,7 +44,10 @@ export ArrowFile,
     box_sampler,
     scaled_distribution_sampler,
     general_sampler,
-    compress_batch_arrow
+    compress_batch_arrow,
+    feasibility_evaluator,
+    objective_evaluator,
+    general_evaluator
 
 include("datasetgen.jl")
 include("csvrecorder.jl")
@@ -52,5 +57,6 @@ include("nn_expression.jl")
 include("metrics.jl")
 include("inconvexhull.jl")
 include("samplers.jl")
+include("model_evaluator.jl")
 
 end
diff --git a/src/arrowrecorder.jl b/src/arrowrecorder.jl
index fed60fa..396ab14 100644
--- a/src/arrowrecorder.jl
+++ b/src/arrowrecorder.jl
@@ -7,14 +7,14 @@ Base.string(::Type{ArrowFile}) = "arrow"
 
 Record optimization problem solution to an Arrow file.
 """
-function record(recorder::Recorder{ArrowFile}, id::UUID; input=false)
+function record(recorder::Recorder{ArrowFile}, id::UUID; input = false)
     _filename = input ? filename_input(recorder) : filename(recorder)
     _filename = _filename * "_$(string(id))." * string(ArrowFile)
     model = recorder.model
 
-    status=JuMP.termination_status(model)
-    primal_stat=JuMP.primal_status(model)
-    dual_stat=JuMP.dual_status(model)
+    status = JuMP.termination_status(model)
+    primal_stat = JuMP.primal_status(model)
+    dual_stat = JuMP.dual_status(model)
 
     primal_values = if in(primal_stat, DECISION_STATUS)
         [[value.(p)] for p in recorder.primal_variables]
@@ -35,25 +35,19 @@ function record(recorder::Recorder{ArrowFile}, id::UUID; input=false)
     end
 
     df = (;
-        id=[id],
-        zip(
-            Symbol.(name.(recorder.primal_variables)),
-            primal_values,
-        )...,
-        zip(
-            Symbol.("dual_" .* name.(recorder.dual_variables)),
-            dual_values,
-        )...,
+        id = [id],
+        zip(Symbol.(name.(recorder.primal_variables)), primal_values)...,
+        zip(Symbol.("dual_" .* name.(recorder.dual_variables)), dual_values)...,
     )
     if !input
         df = merge(
             df,
             (;
-                objective=[objective],
-                time=[JuMP.solve_time(model)],
-                status=[string(status)],
-                primal_status=[string(primal_stat)],
-                dual_status=[string(dual_stat)],
+                objective = [objective],
+                time = [JuMP.solve_time(model)],
+                status = [string(status)],
+                primal_status = [string(primal_stat)],
+                dual_status = [string(dual_stat)],
             ),
         )
     end
@@ -67,13 +61,26 @@ function save(table::NamedTuple, filename::String, ::Type{ArrowFile})
 end
 
 function load(filename::String, ::Type{ArrowFile})
-    return DataFrame(Arrow.Table(filename * "." * string(ArrowFile)))
+    if !occursin(string(ArrowFile), filename)
+        return DataFrame(Arrow.Table(filename * "." * string(ArrowFile)))
+    else
+        return DataFrame(Arrow.Table(filename))
+    end
 end
 
-function compress_batch_arrow(case_file_path::String, case_name::String; keyword_all="output", batch_id::String=string(uuid1()), keyword_any=["_"])
-    iter_files = readdir(case_file_path; join=true)
+function compress_batch_arrow(
+    case_file_path::String,
+    case_name::String;
+    keyword_all = "output",
+    batch_id::String = string(uuid1()),
+    keyword_any = ["_"],
+)
+    iter_files = readdir(case_file_path; join = true)
     file_outs = [
-        file for file in iter_files if occursin(case_name, file) && occursin("arrow", file) && occursin(keyword_all, file) && any(x -> occursin(x, file), keyword_any)
+        file for file in iter_files if occursin(case_name, file) &&
+        occursin("arrow", file) &&
+        occursin(keyword_all, file) &&
+        any(x -> occursin(x, file), keyword_any)
     ]
     output_table = Arrow.Table(file_outs)
     Arrow.write(
diff --git a/src/csvrecorder.jl b/src/csvrecorder.jl
index 080cc29..b209aeb 100644
--- a/src/csvrecorder.jl
+++ b/src/csvrecorder.jl
@@ -7,13 +7,13 @@ Base.string(::Type{CSVFile}) = "csv"
 
 Record optimization problem solution to a CSV file.
 """
-function record(recorder::Recorder{CSVFile}, id::UUID; input=false)
+function record(recorder::Recorder{CSVFile}, id::UUID; input = false)
     _filename = input ? filename_input(recorder) : filename(recorder)
     _filename = _filename * "." * string(CSVFile)
 
     model = recorder.model
-    primal_stat=JuMP.primal_status(model)
-    dual_stat=JuMP.dual_status(model)
+    primal_stat = JuMP.primal_status(model)
+    dual_stat = JuMP.dual_status(model)
 
     if !isfile(_filename)
         open(_filename, "w") do f
@@ -86,11 +86,14 @@ function save(table::NamedTuple, filename::String, ::Type{CSVFile}; kwargs...)
     isappend = isfile(filename)
     mode = isappend ? "append" : "write"
     @info "Saving CSV file to $filename - Mode: $mode"
-    CSV.write(filename, table; append=isappend)
+    CSV.write(filename, table; append = isappend)
     return nothing
 end
 
 function load(filename::String, ::Type{CSVFile})
-    filename = filename * "." * string(CSVFile)
-    return CSV.read(filename, DataFrame)
+    if !occursin(string(CSVFile), filename)
+        return CSV.read(filename * "." * string(CSVFile), DataFrame)
+    else
+        return CSV.read(filename, DataFrame)
+    end
 end
diff --git a/src/datasetgen.jl b/src/datasetgen.jl
index 65140cc..c9f43d2 100644
--- a/src/datasetgen.jl
+++ b/src/datasetgen.jl
@@ -20,7 +20,7 @@ termination_status_filter(status) = in(status, ACCEPTED_TERMINATION_STATUSES)
 primal_status_filter(status) = in(status, DECISION_STATUS)
 dual_status_filter(status) = in(status, DECISION_STATUS)
 
-function filter_fn(model; check_primal=true, check_dual=true)
+function filter_fn(model; check_primal = true, check_dual = true)
     if !termination_status_filter(termination_status(model))
         return false
     elseif check_primal && !primal_status_filter(primal_status(model))
@@ -46,11 +46,11 @@ mutable struct Recorder{T<:FileType}
 
     function Recorder{T}(
         filename::String;
-        filename_input::String=filename * "_input_",
-        primal_variables=[],
-        dual_variables=[],
-        filterfn=filter_fn,
-        model= if length(primal_variables) > 0
+        filename_input::String = filename * "_input_",
+        primal_variables = [],
+        dual_variables = [],
+        filterfn = filter_fn,
+        model = if length(primal_variables) > 0
             owner_model(primal_variables[1])
         elseif length(dual_variables) > 0
             owner_model(dual_variables[1])
@@ -78,10 +78,10 @@ get_filterfn(recorder::Recorder) = recorder.filterfn
 function similar(recorder::Recorder{T}) where {T<:FileType}
     return Recorder{T}(
         filename(recorder);
-        filename_input=filename_input(recorder),
-        primal_variables=get_primal_variables(recorder),
-        dual_variables=get_dual_variables(recorder),
-        filterfn=get_filterfn(recorder),
+        filename_input = filename_input(recorder),
+        primal_variables = get_primal_variables(recorder),
+        dual_variables = get_dual_variables(recorder),
+        filterfn = get_filterfn(recorder),
     )
 end
 
@@ -94,7 +94,7 @@ function set_dual_variable!(recorder::Recorder, p::Vector)
 end
 
 function set_model!(recorder::Recorder)
-    recorder.model= if length(recorder.primal_variables) > 0
+    recorder.model = if length(recorder.primal_variables) > 0
         owner_model(recorder.primal_variables[1])
     elseif length(recorder.dual_variables) > 0
         owner_model(recorder.dual_variables[1])
@@ -128,9 +128,9 @@ struct ProblemIterator{T<:Real} <: AbstractProblemIterator
     function ProblemIterator(
         ids::Vector{UUID},
         pairs::Dict{VariableRef,Vector{T}},
-        early_stop::Function=(args...) -> false,
-        param_type::Type{<:AbstractParameterType}=POIParamaterType,
-        pre_solve_hook::Function=(args...) -> nothing
+        early_stop::Function = (args...) -> false,
+        param_type::Type{<:AbstractParameterType} = POIParamaterType,
+        pre_solve_hook::Function = (args...) -> nothing,
     ) where {T<:Real}
         model = JuMP.owner_model(first(keys(pairs)))
         for (p, val) in pairs
@@ -141,10 +141,11 @@ struct ProblemIterator{T<:Real} <: AbstractProblemIterator
 end
 
 function ProblemIterator(
-    pairs::Dict{VariableRef,Vector{T}}; early_stop::Function=(args...) -> false,
-    pre_solve_hook::Function=(args...) -> nothing,
-    param_type::Type{<:AbstractParameterType}=POIParamaterType,
-    ids = [uuid1() for _ in 1:length(first(values(pairs)))]
+    pairs::Dict{VariableRef,Vector{T}};
+    early_stop::Function = (args...) -> false,
+    pre_solve_hook::Function = (args...) -> nothing,
+    param_type::Type{<:AbstractParameterType} = POIParamaterType,
+    ids = [uuid1() for _ = 1:length(first(values(pairs)))],
 ) where {T<:Real}
     return ProblemIterator(ids, pairs, early_stop, param_type, pre_solve_hook)
 end
@@ -155,10 +156,12 @@ end
 Save optimization problem instances to a file.
 """
 function save(
-    problem_iterator::AbstractProblemIterator, filename::AbstractString, file_type::Type{T}
+    problem_iterator::AbstractProblemIterator,
+    filename::AbstractString,
+    file_type::Type{T},
 ) where {T<:FileType}
-    kys = sort(collect(keys(problem_iterator.pairs)); by=(v) -> index(v).value)
-    df = (; id=problem_iterator.ids,)
+    kys = sort(collect(keys(problem_iterator.pairs)); by = (v) -> index(v).value)
+    df = (; id = problem_iterator.ids,)
     df = merge(df, (; zip(Symbol.(kys), [problem_iterator.pairs[ky] for ky in kys])...))
     save(df, filename, file_type)
     return nothing
@@ -171,8 +174,12 @@ function _dataframe_to_dict(df::DataFrame, parameters::Vector{VariableRef})
             idx = findfirst(parameters) do p
                 name(p) == string(ky)
             end
+            if isnothing(idx)
+                @error("Parameter $ky not found in model")
+                return nothing
+            end
             parameter = parameters[idx]
-            push!(pairs, parameter => df[!,ky])
+            push!(pairs, parameter => df[!, ky])
         end
     end
     return pairs
@@ -182,14 +189,17 @@ function _dataframe_to_dict(df::DataFrame, model_file::AbstractString)
     # Load model
     model = read_from_file(model_file)
     # Retrieve parameters
-    parameters, _ = LearningToOptimize.load_parameters(model)
+    parameters = LearningToOptimize.load_parameters(model)
     return _dataframe_to_dict(df, parameters)
 end
 
-function load(model_file::AbstractString, input_file::AbstractString, ::Type{T}; 
-    batch_size::Union{Nothing, Integer}=nothing,
-    ignore_ids::Vector{UUID}=UUID[],
-    param_type::Type{<:AbstractParameterType}=JuMPParameterType
+function load(
+    model_file::AbstractString,
+    input_file::AbstractString,
+    ::Type{T};
+    batch_size::Union{Nothing,Integer} = nothing,
+    ignore_ids::Vector{UUID} = UUID[],
+    param_type::Type{<:AbstractParameterType} = JuMPParameterType,
 ) where {T<:FileType}
     # Load full set
     df = load(input_file, T)
@@ -206,13 +216,17 @@ function load(model_file::AbstractString, input_file::AbstractString, ::Type{T};
     # No batch
     if isnothing(batch_size)
         pairs = _dataframe_to_dict(df, model_file)
-        return ProblemIterator(pairs; ids=ids, param_type=param_type)
+        return ProblemIterator(pairs; ids = ids, param_type = param_type)
     end
     # Batch
     num_batches = ceil(Int, length(ids) / batch_size)
-    idx_range = (i) -> (i-1)*batch_size+1:min(i*batch_size, length(ids))
-    return (i) -> ProblemIterator(_dataframe_to_dict(df[idx_range(i), :], model_file);
-        ids=ids[idx_range(i)], param_type=param_type), num_batches
+    idx_range = (i) -> (i-1)*batch_size+1:min(i * batch_size, length(ids))
+    return (i) -> ProblemIterator(
+        _dataframe_to_dict(df[idx_range(i), :], model_file);
+        ids = ids[idx_range(i)],
+        param_type = param_type,
+    ),
+    num_batches
 end
 
 """
@@ -237,7 +251,12 @@ end
 
 Update the values of parameters in a JuMP model.
 """
-function update_model!(model::JuMP.Model, pairs::Dict, idx::Integer, param_type::Type{<:AbstractParameterType})
+function update_model!(
+    model::JuMP.Model,
+    pairs::Dict,
+    idx::Integer,
+    param_type::Type{<:AbstractParameterType},
+)
     for (p, val) in pairs
         update_model!(param_type, model, p, val[idx])
     end
@@ -249,7 +268,9 @@ end
 Solve an optimization problem and record the solution.
 """
 function solve_and_record(
-    problem_iterator::ProblemIterator, recorder::Recorder, idx::Integer
+    problem_iterator::ProblemIterator,
+    recorder::Recorder,
+    idx::Integer,
 )
     model = problem_iterator.model
     problem_iterator.pre_solve_hook(model)
@@ -271,7 +292,7 @@ Solve a batch of optimization problems and record the solutions.
 """
 function solve_batch(problem_iterator::AbstractProblemIterator, recorder)
     successfull_solves = 0.0
-    for idx in 1:length(problem_iterator.ids)
+    for idx = 1:length(problem_iterator.ids)
         _success_bool, early_stop_bool = solve_and_record(problem_iterator, recorder, idx)
         if _success_bool == 1
             successfull_solves += 1
diff --git a/src/inconvexhull.jl b/src/inconvexhull.jl
index 6ee4614..b480e9e 100644
--- a/src/inconvexhull.jl
+++ b/src/inconvexhull.jl
@@ -1,15 +1,21 @@
 """
 
-    inconvexhull(training_set::Matrix{Float64}, test_set::Matrix{Float64})
+    inconvexhull(training_set::Matrix{Float64}, test_set::Matrix{Float64}, solver; silent=true, tol=1e-4)
 
 Check if new points are inside the convex hull of the given points. Solves a linear programming problem to check if the points are inside the convex hull.
 """
-function inconvexhull(training_set::Matrix{Float64}, test_set::Matrix{Float64}, solver; silent=true, tol=1e-4)
+function inconvexhull(
+    training_set::Matrix{Float64},
+    test_set::Matrix{Float64},
+    solver;
+    silent = true,
+    tol = 1e-4,
+)
     # Get the number of points and dimensions
     n, d = size(training_set)
     m, d_ = size(test_set)
     @assert d == d_ "The dimensions of the training and test sets must be the same"
-    
+
     # Create the POI model
     model = JuMP.Model(() -> POI.Optimizer(solver()))
     if silent
@@ -22,7 +28,7 @@ function inconvexhull(training_set::Matrix{Float64}, test_set::Matrix{Float64},
 
     # Create POI parameters
     @variable(model, test_set_params[1:d] in MOI.Parameter.(0.0))
-    
+
     # slack variables
     @variable(model, slack[1:d])
     @variable(model, abs_slack[1:d] >= 0)
@@ -36,7 +42,7 @@ function inconvexhull(training_set::Matrix{Float64}, test_set::Matrix{Float64},
     @objective(model, Min, sum(abs_slack))
 
     in_convex_hull = Vector{Bool}(undef, m)
-    for i in 1:m
+    for i = 1:m
         # Set the test set parameters
         MOI.set.(model, POI.ParameterValue(), test_set_params, test_set[i, :])
 
@@ -46,6 +52,6 @@ function inconvexhull(training_set::Matrix{Float64}, test_set::Matrix{Float64},
         # return if the points are inside the convex hull
         in_convex_hull[i] = JuMP.objective_value(model) <= tol
     end
-    
+
     return in_convex_hull
-end
\ No newline at end of file
+end
diff --git a/src/metrics.jl b/src/metrics.jl
index cff131d..8d8984b 100644
--- a/src/metrics.jl
+++ b/src/metrics.jl
@@ -4,4 +4,4 @@ end
 
 function relative_mae(ŷ, y)
     return mean(abs.((ŷ .- y) ./ y))
-end
\ No newline at end of file
+end
diff --git a/src/model_evaluator.jl b/src/model_evaluator.jl
new file mode 100644
index 0000000..e9f44b1
--- /dev/null
+++ b/src/model_evaluator.jl
@@ -0,0 +1,165 @@
+
+"""
+    all_primal_variables(model::JuMP.Model)
+
+Returns all primal variables of a JuMP model.
+"""
+function all_primal_variables(model::JuMP.Model)
+    return sort(
+        setdiff(all_variables(model), load_parameters(model));
+        by = (v) -> index(v).value,
+    )
+end
+
+"""
+    model_outputs(model::Function, problem_iterator::ProblemIterator)
+
+Returns the output of a machine learning model over a ProblemIterator dataset.
+"""
+function model_outputs(model::Function, problem_iterator::ProblemIterator)
+    kys = sort(collect(keys(problem_iterator.pairs)); by = (v) -> index(v).value)
+    input = hcat([problem_iterator.pairs[ky] for ky in kys]...)
+    return model(input)
+end
+
+"""
+    model_outputs(mach::T, problem_iterator::ProblemIterator)
+
+Returns the output of a machine learning model over a ProblemIterator dataset.
+"""
+function model_outputs(mach::T, problem_iterator::ProblemIterator) where {T<:Machine}
+    return predict(
+        mach,
+        DataFrame(Symbol.(keys(problem_iterator.pairs)) .=> values(problem_iterator.pairs)),
+    )
+end
+
+""" 
+    variable_point(vals::Array{T,2}, problem_iterator::ProblemIterator, idx::Int)
+
+Returns a dictionary with the values of the variables and parameters for the idx-th problem instance.
+"""
+function variable_point(
+    vals::Array{T,2},
+    problem_iterator::ProblemIterator,
+    idx::Int,
+) where {T<:Real}
+    # primal variables 
+    variables = all_primal_variables(problem_iterator.model)
+    dict = Dict{VariableRef,Float64}(variables .=> vals[:, idx])
+    # parameters
+    for (v, val) in problem_iterator.pairs
+        dict[v] = val[idx]
+    end
+    return dict
+end
+
+""" 
+    variable_point(vals::Tables.MatrixTable, problem_iterator::ProblemIterator, idx::Int)
+
+Returns a dictionary with the values of the variables and parameters for the idx-th problem instance.
+"""
+function variable_point(
+    vals::Tables.MatrixTable,
+    problem_iterator::ProblemIterator,
+    idx::Int,
+)
+    # primal variables
+    variables = all_primal_variables(problem_iterator.model)
+    dict = Dict{VariableRef,Float64}(
+        variables .=> [vals[Symbol(name(v))][idx] for v in variables],
+    )
+    # parameters
+    for (v, val) in problem_iterator.pairs
+        dict[v] = val[idx]
+    end
+    return dict
+end
+
+function JuMP.MOI.Utilities.distance_to_set(
+    ::MOI.Utilities.ProjectionUpperBoundDistance,
+    val::Float64,
+    set::MOI.Parameter,
+)
+    return set.value - val
+end
+
+
+"""
+    feasibility_evaluator(problem_iterator::ProblemIterator, output)
+
+Feasibility evaluator of a solution over a ProblemIterator dataset.
+PS.: Probably can be made much more efficient.
+"""
+function feasibility_evaluator(problem_iterator::ProblemIterator, output)
+    average_infeasilibity = Array{Float64}(undef, length(problem_iterator.ids))
+    for idx = 1:length(problem_iterator.ids)
+        update_model!(
+            problem_iterator.model,
+            problem_iterator.pairs,
+            idx,
+            problem_iterator.param_type,
+        )
+        dct = primal_feasibility_report(
+            problem_iterator.model,
+            variable_point(output, problem_iterator, idx),
+        )
+        if isempty(dct)
+            average_infeasilibity[idx] = 0.0
+        else
+            average_infeasilibity[idx] = sum(values(dct)) / length(dct)
+        end
+    end
+    return average_infeasilibity
+end
+    
+"""
+    _objective_value(model::JuMP.Model, point::Dict{VariableRef,Float64})
+
+Objective value of a solution over a JuMP model.
+"""
+function _objective_value(model::JuMP.Model, point::Dict{VariableRef,Float64})
+    obj_function = JuMP.objective_function(model)
+    return value((v) -> point[v], obj_function)
+end
+
+"""
+    objective_evaluator(problem_iterator::ProblemIterator, output)
+
+Objective evaluator of a solution over a ProblemIterator dataset.
+PS.: Probably can be made much more efficient.
+"""
+function objective_evaluator(problem_iterator::ProblemIterator, output)
+    average_objective = Array{Float64}(undef, length(problem_iterator.ids))
+    for idx = 1:length(problem_iterator.ids)
+        update_model!(
+            problem_iterator.model,
+            problem_iterator.pairs,
+            idx,
+            problem_iterator.param_type,
+        )
+        average_objective[idx] = _objective_value(
+            problem_iterator.model,
+            variable_point(output, problem_iterator, idx),
+        )
+    end
+    return average_objective
+end
+
+"""
+    general_evaluator(problem_iterator::ProblemIterator, model)
+
+General evaluator of a predictor model over a ProblemIterator dataset. Returns objective, feasibility, 
+inference time and allocated memory.
+"""
+function general_evaluator(problem_iterator::ProblemIterator, model)
+    timed_output = @timed model_outputs(model, problem_iterator)
+    feasibility = feasibility_evaluator(problem_iterator, timed_output.value)
+    objective = objective_evaluator(problem_iterator, timed_output.value)
+    return (;
+        objective = objective,
+        infeasibility = feasibility,
+        time = timed_output.time,
+        bytes = timed_output.bytes,
+    )
+end
diff --git a/src/nn_expression.jl b/src/nn_expression.jl
index 7807508..093b391 100644
--- a/src/nn_expression.jl
+++ b/src/nn_expression.jl
@@ -9,10 +9,8 @@ function NNlib.relu(ex::AffExpr)
     model = owner_model(ex)
     aux = @variable(model, binary = true)
     relu_out = @variable(model, lower_bound = 0.0)
-    @constraint(model, relu_out >= ex * (1-tol))
-    @constraint(model, relu_out <= ex * (1+tol) + big_M * (1 - aux))
+    @constraint(model, relu_out >= ex * (1 - tol))
+    @constraint(model, relu_out <= ex * (1 + tol) + big_M * (1 - aux))
     @constraint(model, relu_out <= big_M * aux)
     return relu_out
 end
-
-
diff --git a/src/samplers.jl b/src/samplers.jl
index 4dceba3..e79a2fe 100644
--- a/src/samplers.jl
+++ b/src/samplers.jl
@@ -11,26 +11,36 @@ The idea is to change the value of one parameter and keep the rest constant.
 function line_sampler(
     original_parameters::Vector{T},
     parameter_indexes::AbstractVector{F},
-    range_p::AbstractVector{T}=1.01:0.01:1.25,
+    range_p::AbstractVector{T} = 1.01:0.01:1.25,
 ) where {T<:Real,F<:Integer}
     parameters = hcat(fill(original_parameters, length(range_p))...)
 
     for parameter_index in parameter_indexes
-        parameters[parameter_index, :] = [original_parameters[parameter_index] * mul for mul in range_p]
+        parameters[parameter_index, :] =
+            [original_parameters[parameter_index] * mul for mul in range_p]
     end
 
     return parameters
 end
 
 function line_sampler(
-    original_parameters::Vector{T},
-    range_p::AbstractVector{T}=1.01:0.01:1.25,
+    _original_parameters::Vector{VariableRef},
+    range_p::AbstractVector{T} = 1.01:0.01:1.25,
 ) where {T<:Real}
-    parameters = zeros(T, length(original_parameters), length(range_p) * (1 + length(original_parameters)))
-    parameters[:, 1:length(range_p)] = line_sampler(original_parameters, 1:length(original_parameters), range_p)
-
-    for parameter_index=1:length(original_parameters)
-        parameters[:, length(range_p) * parameter_index + 1:length(range_p) * (parameter_index + 1)] = line_sampler(original_parameters, [parameter_index], range_p)
+    original_parameters = parameter_value.(_original_parameters)
+    parameters = zeros(
+        T,
+        length(original_parameters),
+        length(range_p) * (1 + length(original_parameters)),
+    )
+    parameters[:, 1:length(range_p)] =
+        line_sampler(original_parameters, 1:length(original_parameters), range_p)
+
+    for parameter_index = 1:length(original_parameters)
+        parameters[
+            :,
+            length(range_p)*parameter_index+1:length(range_p)*(parameter_index+1),
+        ] = line_sampler(original_parameters, [parameter_index], range_p)
     end
 
     return parameters
@@ -38,78 +48,100 @@ end
 
 """
     function box_sampler(
-        original_parameter::T,
-        num_s::F,
-        range_p::AbstractVector{T}=0.8:0.01:1.25,
+        original_parameter::VariableRef,
+        num_s::F;
+        rng::AbstractRNG=Random.GLOBAL_RNG,
+        max_multiplier::T=1.25,
+        min_multiplier::T=0.8,
     ) where {T<:Real,F<:Integer}
     
 Uniformly sample values around the original parameter value over a discrete range inside a box.
 """
 function box_sampler(
-    original_parameter::T,
-    num_s::F,
-    range_p::AbstractVector{T}=0.8:0.01:1.25,
+    original_parameter::VariableRef,
+    num_s::F;
+    rng::AbstractRNG = Random.GLOBAL_RNG,
+    max_multiplier::T = 1.25,
+    min_multiplier::T = 0.8,
 ) where {T<:Real,F<:Integer}
     parameter_samples =
-        original_parameter * rand(range_p, num_s)
+        parameter_value(original_parameter) .*
+        rand(rng, Uniform(min_multiplier, max_multiplier), num_s)
     return parameter_samples
 end
 
 function box_sampler(
-    original_parameters::Vector{T},
-    num_s::F,
-    range_p::AbstractVector{T}=0.8:0.01:1.25,
+    original_parameters::Vector{VariableRef},
+    num_s::F;
+    rng::AbstractRNG = Random.GLOBAL_RNG,
+    max_multiplier::T = 1.25,
+    min_multiplier::T = 0.8,
 ) where {T<:Real,F<:Integer}
-    return vcat([box_sampler(p, num_s, range_p)' for p in original_parameters]...)
+    return vcat(
+        [
+            box_sampler(
+                p,
+                num_s;
+                rng = rng,
+                max_multiplier = max_multiplier,
+                min_multiplier = min_multiplier,
+            )' for p in original_parameters
+        ]...,
+    )
 end
 
 """
     function scaled_distribution_sampler(
-        original_parameters::Vector{T},
+        original_parameters::Vector{VariableRef},
         num_s::F;
         rng::AbstractRNG=Random.GLOBAL_RNG,
         scaler_multiplier::Distribution=Uniform(0.8, 1.25),
         distribution::Distribution=MvLogNormal(fill(-(1.05 .^ 2) ./ 2.0, length(original_parameters)), 1.05)
-    ) where {T<:Real,F<:Integer}
+    ) where {F<:Integer}
 
 Sample from a distribution and scale the parameters by a random value over a uniform distribution.
 """
 function scaled_distribution_sampler(
-    original_parameters::Vector{T},
+    original_parameters::Vector{VariableRef},
     num_s::F;
-    rng::AbstractRNG=Random.GLOBAL_RNG,
-    scaler_multiplier::Distribution=Uniform(0.8, 1.25),
-    distribution::Distribution=MvLogNormal(fill(-(1.05 .^ 2) ./ 2.0, length(original_parameters)), 1.05)
-) where {T<:Real,F<:Integer}
+    rng::AbstractRNG = Random.GLOBAL_RNG,
+    scaler_multiplier::Distribution = Uniform(0.8, 1.25),
+    distribution::Distribution = MvLogNormal(
+        fill(-(1.05 .^ 2) ./ 2.0, length(original_parameters)),
+        1.05,
+    ),
+) where {F<:Integer}
     column_scales = rand(rng, scaler_multiplier, num_s)
     parameter_samples = rand(rng, distribution, num_s)
-    
-    for n in 1:num_s
-        parameter_samples[:, n] = original_parameters .* parameter_samples[:, n] .* column_scales[n]
+
+    for n = 1:num_s
+        parameter_samples[:, n] =
+            parameter_value.(original_parameters) .* parameter_samples[:, n] .*
+            column_scales[n]
     end
     return parameter_samples
 end
 
 """
     function general_sampler(
-        original_parameters::Vector{T};
+        original_parameters::Vector{VariableRef};
         samplers::Vector{Function}=[
             (original_parameters) -> scaled_distribution_sampler(original_parameters, 1000),
             LearningToOptimize.line_sampler, 
             (original_parameters) -> box_sampler(original_parameters, 10),
         ]
-    ) where {T<:Real}
+    )
 
 This function is a general sampler that uses a set of samplers to sample the parameter space.
 """
 function general_sampler(
-    original_parameters::Vector{T};
-    samplers::Vector{Function}=[
+    original_parameters::Vector{VariableRef};
+    samplers::Vector{Function} = [
         (original_parameters) -> scaled_distribution_sampler(original_parameters, 1000),
-        LearningToOptimize.line_sampler, 
+        LearningToOptimize.line_sampler,
         (original_parameters) -> box_sampler(original_parameters, 10),
-    ]
-) where {T<:Real}
+    ],
+)
     return hcat([sampler(original_parameters) for sampler in samplers]...)
 end
 
@@ -119,10 +151,16 @@ end
 Load the parameters from a JuMP model.
 """
 function load_parameters(model::JuMP.Model)
-    cons = constraint_object.([all_constraints(model, VariableRef, MOI.Parameter{Float64}); all_constraints(model, VariableRef, MOI.EqualTo{Float64})])
-    parameters = [cons[i].func for i in 1:length(cons)]
-    vals = [cons[i].set.value for i in 1:length(cons)]
-    return parameters, vals
+    cons =
+        constraint_object.(
+            [
+                all_constraints(model, VariableRef, MOI.Parameter{Float64})
+                all_constraints(model, VariableRef, MOI.EqualTo{Float64})
+            ]
+        )
+    parameters = [cons[i].func for i = 1:length(cons)]
+    parameters = sort(parameters; by = (v) -> index(v).value)
+    return parameters
 end
 
 """
@@ -153,24 +191,19 @@ It saves the sampled parameters to `save_file`.
 """
 function general_sampler(
     file::AbstractString;
-    samplers::Vector{Function}=[
+    samplers::Vector{Function} = [
         (original_parameters) -> scaled_distribution_sampler(original_parameters, 1000),
-        LearningToOptimize.line_sampler, 
+        LearningToOptimize.line_sampler,
         (original_parameters) -> box_sampler(original_parameters, 10),
     ],
-    batch_id::UUID=uuid1(),
-    save_file::AbstractString=split(file, ".mof.json")[1] * "_input_" * string(batch_id),
-    filetype::Type{T}=ArrowFile
+    batch_id::UUID = uuid1(),
+    save_file::AbstractString = split(file, ".mof.json")[1] * "_input_" * string(batch_id),
+    filetype::Type{T} = ArrowFile,
 ) where {T<:FileType}
-    parameters, original_values = load_parameters(file)
-    vals = general_sampler(original_values, samplers=samplers)
-    problem_iterator = ProblemIterator(
-        Dict(parameters .=> [Vector(r) for r in eachrow(vals)]),
-    )
-    save(
-        problem_iterator,
-        save_file,
-        filetype,
-    )
+    parameters = load_parameters(file)
+    vals = general_sampler(parameters, samplers = samplers)
+    problem_iterator =
+        ProblemIterator(Dict(parameters .=> [Vector(r) for r in eachrow(vals)]))
+    save(problem_iterator, save_file, filetype)
     return problem_iterator
-end
\ No newline at end of file
+end
diff --git a/test/datasetgen.jl b/test/datasetgen.jl
index dc3bedf..d56fe6a 100644
--- a/test/datasetgen.jl
+++ b/test/datasetgen.jl
@@ -18,10 +18,12 @@ function test_problem_iterator(path::AbstractString)
         batch_id = string(uuid1())
         @testset "Problem Iterator Builder" begin
             @test_throws AssertionError ProblemIterator(
-                [uuid1() for _ in 1:num_p], Dict(p => collect(1.0:3.0))
+                [uuid1() for _ = 1:num_p],
+                Dict(p => collect(1.0:3.0)),
             )
             @test_throws MethodError ProblemIterator(
-                collect(1.0:3.0), Dict(p => collect(1.0:3.0))
+                collect(1.0:3.0),
+                Dict(p => collect(1.0:3.0)),
             )
             problem_iterator = ProblemIterator(Dict(p => collect(1.0:num_p)))
             file_input = joinpath(path, "test_$(batch_id)_input") # file path
@@ -35,21 +37,21 @@ function test_problem_iterator(path::AbstractString)
         # The recorder
         file_output = joinpath(path, "test_$(batch_id)_output") # file path
         @testset "Recorder Builder" begin
-            @test Recorder{filetype}(file_output; primal_variables=[x]) isa
-                Recorder{filetype}
-            @test Recorder{filetype}(file_output; dual_variables=[cons]) isa
-                Recorder{filetype}
+            @test Recorder{filetype}(file_output; primal_variables = [x]) isa
+                  Recorder{filetype}
+            @test Recorder{filetype}(file_output; dual_variables = [cons]) isa
+                  Recorder{filetype}
         end
-        recorder = Recorder{filetype}(
-            file_output; primal_variables=[x], dual_variables=[cons]
-        )
+        recorder =
+            Recorder{filetype}(file_output; primal_variables = [x], dual_variables = [cons])
 
         # Solve all problems and record solutions
         @testset "early_stop" begin
             file_dual_output = joinpath(path, "test_$(string(uuid1()))_output") # file path
-            recorder_dual = Recorder{filetype}(file_dual_output; dual_variables=[cons])
+            recorder_dual = Recorder{filetype}(file_dual_output; dual_variables = [cons])
             problem_iterator = ProblemIterator(
-                Dict(p => collect(1.0:num_p)); early_stop=(args...) -> true
+                Dict(p => collect(1.0:num_p));
+                early_stop = (args...) -> true,
             )
             successfull_solves = solve_batch(problem_iterator, recorder_dual)
             @test num_p * successfull_solves == 1
@@ -57,10 +59,11 @@ function test_problem_iterator(path::AbstractString)
 
         @testset "pre_solve_hook" begin
             file_dual_output = joinpath(path, "test_$(string(uuid1()))_output") # file path
-            recorder_dual = Recorder{filetype}(file_dual_output; dual_variables=[cons])
+            recorder_dual = Recorder{filetype}(file_dual_output; dual_variables = [cons])
             sum_p = 0
             problem_iterator = ProblemIterator(
-                Dict(p => collect(1.0:num_p)); pre_solve_hook=(args...) -> sum_p += 1
+                Dict(p => collect(1.0:num_p));
+                pre_solve_hook = (args...) -> sum_p += 1,
             )
             successfull_solves = solve_batch(problem_iterator, recorder_dual)
             @test sum_p == num_p
@@ -80,7 +83,7 @@ function test_problem_iterator(path::AbstractString)
                 file_output = file_output * ".$(string(filetype))"
                 @test isfile(file_output)
                 @test length(readdlm(file_output, ',')[:, 1]) ==
-                    num_p * successfull_solves + 1 # 1 from header
+                      num_p * successfull_solves + 1 # 1 from header
                 @test length(readdlm(file_output, ',')[1, :]) == 8
                 rm(file_output)
             else
@@ -115,10 +118,15 @@ function test_problem_iterator(path::AbstractString)
         @objective(model, Min, 2x)
         num_p = 10
         batch_id = string(uuid1())
-        problem_iterator = ProblemIterator(Dict(p => collect(1.0:num_p)); param_type=LearningToOptimize.JuMPParameterType)
+        problem_iterator = ProblemIterator(
+            Dict(p => collect(1.0:num_p));
+            param_type = LearningToOptimize.JuMPParameterType,
+        )
         file_output = joinpath(path, "test_$(batch_id)_output") # file path
         recorder = Recorder{ArrowFile}(
-            file_output; primal_variables=[x], dual_variables=[cons]
+            file_output;
+            primal_variables = [x],
+            dual_variables = [cons],
         )
         successfull_solves = solve_batch(problem_iterator, recorder)
         iter_files = readdir(joinpath(path))
@@ -141,10 +149,15 @@ function test_problem_iterator(path::AbstractString)
         @objective(model, Min, 2x)
         num_p = 10
         batch_id = string(uuid1())
-        problem_iterator = ProblemIterator(Dict(p => collect(1.0:num_p)); param_type=LearningToOptimize.JuMPNLPParameterType)
+        problem_iterator = ProblemIterator(
+            Dict(p => collect(1.0:num_p));
+            param_type = LearningToOptimize.JuMPNLPParameterType,
+        )
         file_output = joinpath(path, "test_$(batch_id)_output") # file path
         recorder = Recorder{ArrowFile}(
-            file_output; primal_variables=[x], dual_variables=[cons]
+            file_output;
+            primal_variables = [x],
+            dual_variables = [cons],
         )
         successfull_solves = solve_batch(problem_iterator, recorder)
         iter_files = readdir(joinpath(path))
@@ -161,8 +174,12 @@ function test_problem_iterator(path::AbstractString)
     end
 end
 
-function test_load(model_file::AbstractString, input_file::AbstractString, ::Type{T}, ids::Vector{UUID};
-    batch_size::Integer=32
+function test_load(
+    model_file::AbstractString,
+    input_file::AbstractString,
+    ::Type{T},
+    ids::Vector{UUID};
+    batch_size::Integer = 32,
 ) where {T<:LearningToOptimize.FileType}
     # Test Load full set 
     problem_iterator = LearningToOptimize.load(model_file, input_file, T)
@@ -170,30 +187,67 @@ function test_load(model_file::AbstractString, input_file::AbstractString, ::Typ
     @test length(problem_iterator.ids) == length(ids)
     # Test load only half of the ids
     num_ids_ignored = floor(Int, length(ids) / 2)
-    problem_iterator =  LearningToOptimize.load(model_file, input_file, T; ignore_ids=ids[1:num_ids_ignored])
+    problem_iterator = LearningToOptimize.load(
+        model_file,
+        input_file,
+        T;
+        ignore_ids = ids[1:num_ids_ignored],
+    )
     @test length(problem_iterator.ids) == length(ids) - num_ids_ignored
     # Test warning all ids to be ignored
-    problem_iterator =  LearningToOptimize.load(model_file, input_file, T; ignore_ids=ids)
+    problem_iterator = LearningToOptimize.load(model_file, input_file, T; ignore_ids = ids)
     @test isnothing(problem_iterator)
     # Test Load batch of problem iterators
-    problem_iterator_factory, num_batches =  LearningToOptimize.load(model_file, input_file, T; batch_size=batch_size)
+    problem_iterator_factory, num_batches =
+        LearningToOptimize.load(model_file, input_file, T; batch_size = batch_size)
     @test num_batches == ceil(Int, length(ids) / batch_size)
     @test problem_iterator_factory(1) isa LearningToOptimize.AbstractProblemIterator
     return nothing
 end
 
-function test_compress_batch_arrow(case_file_path::AbstractString=mktempdir(), case_name::AbstractString="test"; keyword::AbstractString="output")
+function test_compress_batch_arrow(
+    case_file_path::AbstractString = mktempdir(),
+    case_name::AbstractString = "test";
+    keyword::AbstractString = "output",
+)
     random_data = rand(10, 10)
-    col_names = ["col_$(i)" for i in 1:10]
-    batch_ids = [string(uuid1()) for _ in 1:10]
-    dfs = [DataFrame(random_data[i:i, :], col_names) for i in 1:10]
-    for i in 1:10
-        Arrow.write(joinpath(case_file_path, "$(case_name)_$(keyword)_$(batch_ids[i]).arrow"), dfs[i])
+    col_names = ["col_$(i)" for i = 1:10]
+    batch_ids = [string(uuid1()) for _ = 1:10]
+    dfs = [DataFrame(random_data[i:i, :], col_names) for i = 1:10]
+    for i = 1:10
+        Arrow.write(
+            joinpath(case_file_path, "$(case_name)_$(keyword)_$(batch_ids[i]).arrow"),
+            dfs[i],
+        )
     end
-    @test length([file for file in readdir(case_file_path; join=true) if occursin(case_name, file) && occursin("arrow", file) && occursin(keyword, file) && any(x -> occursin(x, file), batch_ids)]) == 10
+    @test length([
+        file for
+        file in readdir(case_file_path; join = true) if occursin(case_name, file) &&
+        occursin("arrow", file) &&
+        occursin(keyword, file) &&
+        any(x -> occursin(x, file), batch_ids)
+    ]) == 10
     batch_id = string(uuid1())
-    LearningToOptimize.compress_batch_arrow(case_file_path, case_name; keyword_all=keyword, batch_id=batch_id, keyword_any=batch_ids)
-    @test length([file for file in readdir(case_file_path; join=true) if occursin(case_name, file) && occursin("arrow", file) && occursin(keyword, file) && any(x -> occursin(x, file), batch_ids)]) == 0
-    @test length([file for file in readdir(case_file_path; join=true) if occursin(case_name, file) && occursin("arrow", file) && occursin(keyword, file) && occursin(batch_id, file)]) == 1
+    LearningToOptimize.compress_batch_arrow(
+        case_file_path,
+        case_name;
+        keyword_all = keyword,
+        batch_id = batch_id,
+        keyword_any = batch_ids,
+    )
+    @test length([
+        file for
+        file in readdir(case_file_path; join = true) if occursin(case_name, file) &&
+        occursin("arrow", file) &&
+        occursin(keyword, file) &&
+        any(x -> occursin(x, file), batch_ids)
+    ]) == 0
+    @test length([
+        file for
+        file in readdir(case_file_path; join = true) if occursin(case_name, file) &&
+        occursin("arrow", file) &&
+        occursin(keyword, file) &&
+        occursin(batch_id, file)
+    ]) == 1
     return nothing
 end
diff --git a/test/inconvexhull.jl b/test/inconvexhull.jl
index 14e6445..794ffc8 100644
--- a/test/inconvexhull.jl
+++ b/test/inconvexhull.jl
@@ -6,12 +6,13 @@ Test the inconvexhull function: inconvexhull(training_set::Matrix{Float64}, test
 function test_inconvexhull()
     @testset "inconvexhull" begin
         # Create the training set
-        training_set = [0. 0; 1 0; 0 1; 1 1]
-        
+        training_set = [0.0 0; 1 0; 0 1; 1 1]
+
         # Create the test set
         test_set = [0.5 0.5; -0.5 0.5; 0.5 -0.5; 0.0 0.5; 2.0 0.5]
-        
+
         # Test the inconvexhull function
-        @test inconvexhull(training_set, test_set, HiGHS.Optimizer) == [true, false, false, true, false]
+        @test inconvexhull(training_set, test_set, HiGHS.Optimizer) ==
+              [true, false, false, true, false]
     end
 end
diff --git a/test/nn_expression.jl b/test/nn_expression.jl
index 5b8b2e4..0a2d577 100644
--- a/test/nn_expression.jl
+++ b/test/nn_expression.jl
@@ -4,10 +4,10 @@
 Tests running a jump model with a flux expression.
 """
 function test_flux_jump_basic()
-    for i in 1:10
+    for i = 1:10
         model = JuMP.Model(HiGHS.Optimizer)
 
-        @variable(model, x[i = 1:3]>= 2.3)
+        @variable(model, x[i = 1:3] >= 2.3)
 
         flux_model = Chain(Dense(3, 3, relu), Dense(3, 1))
 
@@ -22,9 +22,9 @@ function test_flux_jump_basic()
 
         @test termination_status(model) === OPTIMAL
         if flux_model(value.(x))[1] <= 1.0
-            @test isapprox(flux_model(value.(x))[1], value(ex); atol=0.01)
+            @test isapprox(flux_model(value.(x))[1], value(ex); atol = 0.01)
         else
-            @test isapprox(flux_model(value.(x))[1], value(ex); rtol=0.001)
+            @test isapprox(flux_model(value.(x))[1], value(ex); rtol = 0.001)
         end
     end
 end
diff --git a/test/runtests.jl b/test/runtests.jl
index 5768079..b0ae0b8 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -44,11 +44,13 @@ include(joinpath(test_dir, "samplers.jl"))
         test_compress_batch_arrow(path)
         model_file = "pglib_opf_case5_pjm_DCPPowerModel_POI_load.mof.json"
         @testset "Samplers saving on $filetype" for filetype in [ArrowFile, CSVFile]
-            file_in, ids = test_general_sampler_file(model_file; cache_dir=path, filetype=filetype)
+            file_in, ids =
+                test_general_sampler_file(model_file; cache_dir = path, filetype = filetype)
             test_load(model_file, file_in, filetype, ids)
         end
         test_problem_iterator(path)
-        file_in, file_out = test_pglib_datasetgen(path, "pglib_opf_case5_pjm", 20)
-        test_flux_forecaster(file_in, file_out)
+        file_in, file_out, problem_iterator =
+            test_pglib_datasetgen(path, "pglib_opf_case5_pjm", 20)
+        test_flux_forecaster(problem_iterator, file_in, file_out)
     end
 end
diff --git a/test/samplers.jl b/test/samplers.jl
index cb09848..3c82ad9 100644
--- a/test/samplers.jl
+++ b/test/samplers.jl
@@ -1,53 +1,78 @@
-function test_line_sampler(; num_p=10, range_p = 1:0.01:1.1)
+function test_line_sampler(; num_p = 10, range_p = 1:0.01:1.1)
     original_parameter = rand(num_p)
+    model = JuMP.Model()
+    @variable(model, 0 <= x <= 1)
+    @variable(model, p[1:num_p] in MOI.Parameter.(original_parameter))
+    @constraint(model, cons, x + sum(p) >= 3)
+    @objective(model, Min, 2x)
+
     for parameter_index = 1:num_p
-        parameters = LearningToOptimize.line_sampler(
-            original_parameter,
-            [parameter_index],
-            range_p,
-        )
+        parameters =
+            LearningToOptimize.line_sampler(original_parameter, [parameter_index], range_p)
         @test parameters[parameter_index, 1] == original_parameter[parameter_index]
-        @test parameters[parameter_index, :] == [original_parameter[parameter_index] * mul for mul in range_p]
+        @test parameters[parameter_index, :] ==
+              [original_parameter[parameter_index] * mul for mul in range_p]
     end
-    parameters = LearningToOptimize.line_sampler(
-        original_parameter,
-        range_p,
-    )
+    parameters = LearningToOptimize.line_sampler(p, range_p)
     @test size(parameters) == (10, length(range_p) * (1 + num_p))
     return nothing
 end
 
-function test_box_sampler(; num_p=10, num_s=5, max_multiplier=3.0, min_multiplier=0.0, step_multiplier=0.1)
+function test_box_sampler(;
+    num_p = 10,
+    num_s = 5,
+    max_multiplier = 3.0,
+    min_multiplier = 0.0,
+)
     original_parameter = rand(num_p)
-    parameters = box_sampler(original_parameter, num_s, min_multiplier:step_multiplier:max_multiplier)
+    model = JuMP.Model()
+    @variable(model, 0 <= x <= 1)
+    @variable(model, p[1:num_p] in MOI.Parameter.(original_parameter))
+    @constraint(model, cons, x + sum(p) >= 3)
+    @objective(model, Min, 2x)
+
+    parameters = box_sampler(
+        p,
+        num_s;
+        max_multiplier = max_multiplier,
+        min_multiplier = min_multiplier,
+    )
     @test size(parameters) == (num_p, num_s)
     @test all(parameters .>= original_parameter * min_multiplier)
     @test all(parameters .<= original_parameter * max_multiplier)
     return nothing
 end
 
-function test_general_sampler(; num_p=10, num_s=5, range_p=1.01:0.01:1.25)
+function test_general_sampler(; num_p = 10, num_s = 5, range_p = 1.01:0.01:1.25)
     original_parameter = rand(num_p)
+    model = JuMP.Model()
+    @variable(model, 0 <= x <= 1)
+    @variable(model, p[1:num_p] in MOI.Parameter.(original_parameter))
+    @constraint(model, cons, x + sum(p) >= 3)
+    @objective(model, Min, 2x)
+
     parameters = general_sampler(
-        original_parameter;
-        samplers=[
-            (original_parameters) -> scaled_distribution_sampler(original_parameters, num_s),
-            line_sampler, 
+        p;
+        samplers = [
+            (original_parameters) ->
+                scaled_distribution_sampler(original_parameters, num_s),
+            line_sampler,
             (original_parameters) -> box_sampler(original_parameters, num_s),
-        ]
+        ],
     )
     @test size(parameters) == (num_p, 2 * num_s + length(range_p) * (1 + num_p))
     return nothing
 end
 
-function test_load_parameters_model(;num_p=10, num_v=5)
+function test_load_parameters_model(; num_p = 10, num_v = 5)
     model = JuMP.Model()
     @variable(model, 0 <= x[1:num_v] <= 1)
     @variable(model, p[1:num_p] in MOI.Parameter.(1.0))
     @constraint(model, cons, sum(x) + sum(p) >= 3)
     @objective(model, Min, 2x)
 
-    parameters, vals = LearningToOptimize.load_parameters(model)
+    parameters = LearningToOptimize.load_parameters(model)
+    vals = parameter_value.(parameters)
     @test length(parameters) == num_p
     @test length(vals) == num_p
     @test all(vals .== 1.0)
@@ -56,33 +81,40 @@ function test_load_parameters_model(;num_p=10, num_v=5)
 end
 
 function test_load_parameters()
-    file="pglib_opf_case5_pjm_DCPPowerModel_POI_load.mof.json"
-    parameters, vals = LearningToOptimize.load_parameters(file)
+    file = "pglib_opf_case5_pjm_DCPPowerModel_POI_load.mof.json"
+    parameters = LearningToOptimize.load_parameters(file)
+    vals = parameter_value.(parameters)
     @test length(parameters) == 3
     @test length(vals) == 3
     @test all(vals .== 1.0)
     return nothing
 end
 
-function test_general_sampler_file(file::AbstractString="pglib_opf_case5_pjm_DCPPowerModel_POI_load.mof.json"; 
-    num_s=5, range_p=1.01:0.01:1.25, 
-    cache_dir=mktempdir(),
-    batch_id=uuid1(),
-    save_file = joinpath(cache_dir, split(split(file, ".mof.json")[1], "/")[end] * "_input_" * string(batch_id)),
-    filetype=CSVFile
+function test_general_sampler_file(
+    file::AbstractString = "pglib_opf_case5_pjm_DCPPowerModel_POI_load.mof.json";
+    num_s = 5,
+    range_p = 1.01:0.01:1.25,
+    cache_dir = mktempdir(),
+    batch_id = uuid1(),
+    save_file = joinpath(
+        cache_dir,
+        split(split(file, ".mof.json")[1], "/")[end] * "_input_" * string(batch_id),
+    ),
+    filetype = CSVFile,
 )
-    _, vals = LearningToOptimize.load_parameters(file)
-    num_p=length(vals)
+    parameters = LearningToOptimize.load_parameters(file)
+    num_p = length(parameters)
     problem_iterator = general_sampler(
         file;
-        samplers=[
-            (original_parameters) -> scaled_distribution_sampler(original_parameters, num_s),
+        samplers = [
+            (original_parameters) ->
+                scaled_distribution_sampler(original_parameters, num_s),
             (original_parameters) -> line_sampler(original_parameters, range_p),
             (original_parameters) -> box_sampler(original_parameters, num_s),
         ],
-        save_file=save_file,
-        batch_id=batch_id,
-        filetype=filetype
+        save_file = save_file,
+        batch_id = batch_id,
+        filetype = filetype,
     )
     @test length(problem_iterator.ids) == 2 * num_s + length(range_p) * (1 + num_p)
     @test length(problem_iterator.pairs) == num_p
diff --git a/test/test_flux_forecaster.jl b/test/test_flux_forecaster.jl
index c1d8f8c..9fb4143 100644
--- a/test/test_flux_forecaster.jl
+++ b/test/test_flux_forecaster.jl
@@ -1,37 +1,52 @@
 """
-    test_flux_forecaster(file_in::AbstractString, file_out::AbstractString)
+    test_flux_forecaster(problem_iterator::ProblemIterator, file_in::AbstractString, file_out::AbstractString)
 
 Test the Flux.jl forecaster using MLJ and MLJFlux to train the neural network.
 """
-function test_flux_forecaster(file_in::AbstractString, file_out::AbstractString)
+function test_flux_forecaster(
+    problem_iterator::ProblemIterator,
+    file_in::AbstractString,
+    file_out::AbstractString,
+)
     @testset "Flux.jl & MLJ.jl" begin
-        @test sprint(show, FullyConnected(1, [1], 1)) == "FullyConnected(\n     Layers: PairwiseFusion(vcat, Dense(1 => 1, relu), Dense(2 => 1)),\n     Pass-Through: Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}[Dense(1 => 1)]\n)\n"
+        @test sprint(show, FullyConnected(1, [1], 1)) ==
+              "FullyConnected(\n     Layers: PairwiseFusion(vcat, Dense(1 => 1, relu), Dense(2 => 1)),\n     Pass-Through: Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}[Dense(1 => 1)]\n)\n"
+
+        # problem_iterator = LearningToOptimize.load(model_file, file_in, CSVFile)
 
         # read input and output data
         input_data = CSV.read(file_in, DataFrame)
         output_data = CSV.read(file_out, DataFrame)
 
         # Separate input and output variables
-        output_variables = output_data[!, Not([:id, :status, :primal_status, :dual_status])]
-        input_features = innerjoin(input_data, output_data[!, [:id]]; on=:id)[!, Not(:id)] # just use success solves
+        output_variables = output_data[
+            !,
+            Not([:id, :status, :primal_status, :dual_status, :objective, :time]),
+        ]
+        input_features = innerjoin(input_data, output_data[!, [:id]]; on = :id)[!, Not(:id)] # just use success solves
+
+        # Order columns by index
+        # input_features = input_features[:, sort(names(input_features), by=x->parse(Int, split(x, "_")[end]))]
 
         # Define model
         model = MultitargetNeuralNetworkRegressor(;
-            builder=FullyConnectedBuilder([64, 32]),
-            rng=123,
-            epochs=20,
-            optimiser=ConvexRule(
-                Optimisers.Adam()
-            ),
+            builder = FullyConnectedBuilder([64, 32]),
+            rng = 123,
+            epochs = 20,
+            optimiser = ConvexRule(Optimisers.Adam()),
         )
 
         # Define the machine
         mach = machine(model, input_features, output_variables)
-        fit!(mach; verbosity=2)
+        fit!(mach; verbosity = 2)
 
-        # Make predictions
-        predictions = predict(mach, input_features)
-        @test predictions isa Tables.MatrixTable
+        # Make predictions and evaluate
+        evaluation = general_evaluator(problem_iterator, mach)
+        # predictions = predict(mach, input_features)
+        @test evaluation.objective isa Array
+        @test evaluation.infeasibility isa Array
+        @test evaluation.time isa Real
+        @test evaluation.bytes isa Int
 
         # Delete the files
         rm(file_in)
@@ -40,7 +55,7 @@ function test_flux_forecaster(file_in::AbstractString, file_out::AbstractString)
 end
 
 # Test the Flux.jl forecaster outside MLJ.jl
-function test_fully_connected(;num_samples::Int=100, num_features::Int=10)
+function test_fully_connected(; num_samples::Int = 100, num_features::Int = 10)
     X = rand(num_features, num_samples)
     y = rand(1, num_samples)