diff --git a/.gitignore b/.gitignore
index 7f361e7..7d1a06e 100644
--- a/.gitignore
+++ b/.gitignore
@@ -12,9 +12,16 @@ Manifest.toml
 *.out
 *.sbatch
 examples/unitcommitment/app/*
+*/app/*
 *.edu
-examples/unitcommitment/wandb/*
+*.yaml
 *.png
 *.jls
 *.jlso
 *.jld2
+*.txt
+*.json
+*.log
+*.wandb
+*latest-run
+*.html
diff --git a/README.md b/README.md
index 3ccb141..d0cdd87 100644
--- a/README.md
+++ b/README.md
@@ -134,7 +134,7 @@ Flux.train!(loss, Flux.params(model), [(input_features, output_variables)], opti
 predictions = model(input_features)
 ```
 
-## Comming Soon
+## Coming Soon
 
 Future features:
  - ML objectives that penalize infeasible predictions;
diff --git a/examples/powermodels/HuggingFaceDatasets.jl b/examples/powermodels/HuggingFaceDatasets.jl
index 2133064..a3d3ade 100644
--- a/examples/powermodels/HuggingFaceDatasets.jl
+++ b/examples/powermodels/HuggingFaceDatasets.jl
@@ -31,7 +31,7 @@ cache_dir="./examples/powermodels/data/"
 organization = "L2O"
 dataset = "pglib_opf_solves"
 case_name = "pglib_opf_case300_ieee"
-formulation = "DCPPowerModel"
+formulation = "SOCWRConicPowerModel" # ACPPowerModel SOCWRConicPowerModel DCPPowerModel
 io_type = "input"
 download_dataset(organization, dataset, case_name, io_type; cache_dir=cache_dir)
 
diff --git a/examples/powermodels/data_split.jl b/examples/powermodels/data_split.jl
new file mode 100644
index 0000000..a215385
--- /dev/null
+++ b/examples/powermodels/data_split.jl
@@ -0,0 +1,95 @@
+####################################################
+############## PowerModels Data Split ##############
+####################################################
+import Pkg
+Pkg.activate(dirname(dirname(@__DIR__)))
+
+using Distributed
+using Random
+
+##############
+# Load Packages everywhere
+##############
+
+@everywhere import Pkg
+@everywhere Pkg.activate(dirname(dirname(@__DIR__)))
+@everywhere Pkg.instantiate()
+@everywhere using DataFrames
+@everywhere using L2O
+@everywhere using Gurobi
+@everywhere using Arrow
+@everywhere using MLUtils
+
+##############
+# Parameters
+##############
+case_name = "pglib_opf_case300_ieee" # pglib_opf_case300_ieee # pglib_opf_case5_pjm
+filetype = ArrowFile # ArrowFile # CSVFile
+path_dataset = joinpath(dirname(@__FILE__), "data")
+case_file_path = joinpath(path_dataset, case_name)
+case_file_path_input = joinpath(case_file_path, "input")
+
+mkpath(joinpath(case_file_path_input, "train"))
+mkpath(joinpath(case_file_path_input, "test"))
+
+##############
+# Load Data
+##############
+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)
+]
+batch_ids = [split(split(file, "_")[end], ".")[1] for file in file_ins]
+
+# Load input and output data tables
+if filetype === ArrowFile
+    input_table_train = Arrow.Table(file_ins)
+else
+    input_table_train = CSV.read(file_ins[train_idx], DataFrame)
+end
+
+# Convert to dataframes
+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_table = input_data[train_idx, :]
+test_table = input_data[test_idx, :]
+
+batch_size = 10
+
+num_batches = ceil(Int, length(test_idx) / batch_size)
+
+##############
+# Check Convex-Hull
+##############
+
+@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))
+    batch = test_table[idx_range, :]
+    inhull[idx_range] = inconvexhull(Matrix(train_table[!, Not(:id)]), Matrix(batch[!, Not(:id)]), Gurobi.Optimizer)
+    @info "Batch $i of $num_batches done"
+end
+
+test_table.in_train_convex_hull = inhull
+
+##############
+# Save Files
+##############
+
+# 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)
+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
diff --git a/examples/powermodels/flux_forecaster_script.jl b/examples/powermodels/flux_forecaster_script.jl
deleted file mode 100644
index 9594b74..0000000
--- a/examples/powermodels/flux_forecaster_script.jl
+++ /dev/null
@@ -1,97 +0,0 @@
-using TestEnv
-TestEnv.activate()
-
-using Arrow
-using CSV
-using MLJFlux
-using Flux
-using MLJ
-using DataFrames
-using PowerModels
-using L2O
-
-# Paths
-case_name = "pglib_opf_case300_ieee" # pglib_opf_case300_ieee # pglib_opf_case5_pjm
-network_formulation = DCPPowerModel # SOCWRConicPowerModel # DCPPowerModel
-filetype = ArrowFile # ArrowFile # CSVFile
-path_dataset = joinpath(pwd(), "examples", "powermodels", "data")
-case_file_path = joinpath(path_dataset, case_name)
-case_file_path_output = joinpath(case_file_path, "output", string(network_formulation))
-case_file_path_input = joinpath(case_file_path, "input")
-
-# Load input and output data tables
-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)
-]
-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)
-]
-batch_ids = [split(split(file, "_")[end], ".")[1] for file in file_ins]
-
-# Load input and output data tables
-train_idx = collect(1:floor(Int, length(file_ins) * 0.5))
-test_idx = setdiff(1:length(file_ins), train_idx)
-
-if filetype === ArrowFile
-    input_table_train = Arrow.Table(file_ins[train_idx])
-    output_table_train = Arrow.Table(file_outs[train_idx])
-
-    input_table_test = Arrow.Table(file_ins[test_idx])
-    output_table_test = Arrow.Table(file_outs[test_idx])
-else
-    input_table_train = CSV.read(file_ins[train_idx], DataFrame)
-    output_table_train = CSV.read(file_outs[train_idx], DataFrame)
-
-    input_table_test = CSV.read(file_ins[test_idx], DataFrame)
-    output_table_test = CSV.read(file_outs[test_idx], DataFrame)
-end
-
-# Convert to dataframes
-input_data_train = DataFrame(input_table_train)
-output_data_train = DataFrame(output_table_train)
-
-input_data_test = DataFrame(input_table_test)
-output_data_test = DataFrame(output_table_test)
-
-# Separate input and output variables
-output_variables_train = output_data_train[!, Not(:id)]
-input_features_train = innerjoin(input_data_train, output_data_train[!, [:id]]; on=:id)[
-    !, Not(:id)
-] # just use success solves
-
-num_loads = floor(Int, size(input_features_train, 2) / 2)
-total_volume = [
-    sum(
-        sqrt(input_features_train[i, l]^2 + input_features_train[i, l + num_loads]^2) for
-        l in 1:num_loads
-    ) for i in 1:size(input_features_train, 1)
-]
-
-output_variables_test = output_data_test[!, Not(:id)]
-input_features_test = innerjoin(input_data_test, output_data_test[!, [:id]]; on=:id)[
-    !, Not(:id)
-] # just use success solves
-
-# Define model
-model = MultitargetNeuralNetworkRegressor(;
-    builder=FullyConnectedBuilder([64, 32]),
-    rng=123,
-    epochs=20,
-    optimiser=ConvexRule(
-        Flux.Optimise.Adam(0.001, (0.9, 0.999), 1.0e-8, IdDict{Any,Any}())
-    ),
-)
-
-# Define the machine
-mach = machine(model, input_features_train, output_variables_train)
-fit!(mach; verbosity=2)
-
-# Make predictions
-predictions = predict(mach, input_features)
-
-# Calculate the error
-error = Flux.mse(predictions, output_variables_test)
diff --git a/examples/powermodels/jls2jld2.jl b/examples/powermodels/jls2jld2.jl
new file mode 100644
index 0000000..eb870ff
--- /dev/null
+++ b/examples/powermodels/jls2jld2.jl
@@ -0,0 +1,84 @@
+using Arrow
+using CSV
+using MLJFlux
+using MLUtils
+using Flux
+using MLJ
+using CUDA
+using DataFrames
+using PowerModels
+using L2O
+using Random
+using JLD2
+using Wandb, Dates, Logging
+
+include(joinpath(dirname(dirname(@__FILE__)), "training_utils.jl")) # include("../training_utils.jl")
+
+##############
+# Parameters
+##############
+case_name = ARGS[1] # case_name="pglib_opf_case300_ieee" # pglib_opf_case5_pjm
+network_formulation = ARGS[2] # network_formulation=ACPPowerModel SOCWRConicPowerModel DCPPowerModel
+icnn = parse(Bool, ARGS[3]) # icnn=true # false
+filetype = ArrowFile # ArrowFile # CSVFile
+layers = [512, 256, 64] # [256, 64, 32]
+path_dataset = joinpath(dirname(@__FILE__), "data")
+case_file_path = joinpath(path_dataset, case_name)
+case_file_path_output = joinpath(case_file_path, "output", string(network_formulation))
+case_file_path_input = joinpath(case_file_path, "input", "train")
+save_file = if icnn
+    "$(case_name)_$(network_formulation)_$(replace(string(layers), ", " => "_"))_icnn"
+else
+    "$(case_name)_$(network_formulation)_$(replace(string(layers), ", " => "_"))_dnn"
+end
+
+##############
+# Load Data
+##############
+
+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)
+]
+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)
+]
+# batch_ids = [split(split(file, "_")[end], ".")[1] for file in file_ins]
+
+# Load input and output data tables
+if filetype === ArrowFile
+    input_table_train = Arrow.Table(file_ins)
+    output_table_train = Arrow.Table(file_outs)
+else
+    input_table_train = CSV.read(file_ins[train_idx], DataFrame)
+    output_table_train = CSV.read(file_outs[train_idx], DataFrame)
+end
+
+# Convert to dataframes
+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, :]
+
+# match
+train_table = innerjoin(input_data, output_data[!, [:id, :operational_cost]]; on=:id)
+
+input_features = names(train_table[!, Not([:id, :operational_cost])])
+
+##############
+# Load JLS Model
+##############
+num = 3
+model_dir = joinpath(dirname(@__FILE__), "models")
+mach = machine(joinpath(model_dir, save_file * "$num.jls"))
+
+##############
+# Save JLD2 Model
+##############
+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
diff --git a/examples/powermodels/powermodels_nn_testing.jl b/examples/powermodels/powermodels_nn_testing.jl
new file mode 100644
index 0000000..0832e3a
--- /dev/null
+++ b/examples/powermodels/powermodels_nn_testing.jl
@@ -0,0 +1,132 @@
+####################################################
+############ PowerModels Model Testing #############
+####################################################
+
+using Arrow
+using CSV
+using Flux
+using DataFrames
+using PowerModels
+using L2O
+using JLD2
+using Statistics
+
+##############
+# Parameters
+##############
+case_name = ARGS[1] # case_name="pglib_opf_case300_ieee" # pglib_opf_case5_pjm
+network_formulation = ARGS[2] # network_formulation=ACPPowerModel SOCWRConicPowerModel DCPPowerModel
+filetype = ArrowFile # ArrowFile # CSVFile
+path_dataset = joinpath(dirname(@__FILE__), "data")
+case_file_path = joinpath(path_dataset, case_name)
+case_file_path_output = joinpath(case_file_path, "output", string(network_formulation))
+case_file_path_input = joinpath(case_file_path, "input", "test")
+save_file = "$(case_name)_$(network_formulation)"
+
+
+##############
+# Load Data
+##############
+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)
+]
+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)
+]
+
+# Load input and output data tables
+if filetype === ArrowFile
+    input_table_train = Arrow.Table(file_ins)
+    output_table_train = Arrow.Table(file_outs)
+else
+    input_table_train = CSV.read(file_ins[train_idx], DataFrame)
+    output_table_train = CSV.read(file_outs[train_idx], DataFrame)
+end
+
+# Convert to dataframes
+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, :]
+
+# match
+test_table = innerjoin(input_data, output_data[!, [:id, :operational_cost]]; on=:id)
+
+##############
+# Load DNN Approximator
+##############
+
+model_dir = joinpath(dirname(@__FILE__), "models")
+
+models_list = readdir(model_dir)
+models_list = filter(x -> occursin(".jld2", x), models_list)
+models_list = filter(x -> occursin(string(network_formulation), x), models_list)
+
+flux_models = Array{Chain}(undef, length(models_list))
+layers = Array{Vector}(undef, length(models_list))
+type = Array{String}(undef, length(models_list))
+global input_features = 0
+for (i, model) in enumerate(models_list)
+    type[i] = occursin("icnn", model) ? "ICNN" : "DNN"
+    model_save = JLD2.load(joinpath(model_dir, model))
+    model_state = model_save["model_state"]
+    global input_features = model_save["input_features"]
+    layers[i] = model_save["layers"]
+    input_size = length(input_features)
+    flux_model = Chain(FullyConnected(input_size, layers[i], 1))
+    Flux.loadmodel!(flux_model, model_state)
+    flux_models[i] = flux_model
+end
+
+##############
+# Prepare Data
+##############
+
+X = Float32.(Matrix(test_table[!, input_features]))
+y = Float32.(Matrix(test_table[!, [:operational_cost]]))
+
+##############
+# Predict
+##############
+mae_convex_hull = Array{Float32}(undef, length(flux_models))
+mae_out_convex_hull = Array{Float32}(undef, length(flux_models))
+worst_case_convex_hull = Array{Float32}(undef, length(flux_models))
+worst_case_out_convex_hull = Array{Float32}(undef, length(flux_models))
+std_convex_hull = Array{Float32}(undef, length(flux_models))
+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]
+    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))
+    worst_case_out_convex_hull[i] = maximum(abs.(error_vec_out_chull))
+    std_convex_hull[i] = std(error_vec_in_chull)
+    std_out_convex_hull[i] = std(error_vec_out_chull)
+end
+
+# table
+results = DataFrame(
+    type = type,
+    layers = layers,
+    mae_convex_hull = mae_convex_hull,
+    mae_out_convex_hull = mae_out_convex_hull,
+    worst_case_convex_hull = worst_case_convex_hull,
+    worst_case_out_convex_hull = worst_case_out_convex_hull,
+    std_convex_hull = std_convex_hull,
+    std_out_convex_hull = std_out_convex_hull,
+)
+
+# save
+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
diff --git a/examples/powermodels/powermodels_nn_training.jl b/examples/powermodels/powermodels_nn_training.jl
new file mode 100644
index 0000000..94a440a
--- /dev/null
+++ b/examples/powermodels/powermodels_nn_training.jl
@@ -0,0 +1,176 @@
+####################################################
+############ PowerModels Model Training ############
+####################################################
+
+using Arrow
+using CSV
+using MLJFlux
+using MLUtils
+using MLJBase
+using Flux
+using MLJ
+using CUDA
+using DataFrames
+using PowerModels
+using L2O
+using Random
+using JLD2
+using Wandb, Dates, Logging
+
+include(joinpath(dirname(dirname(@__FILE__)), "training_utils.jl")) # include("../training_utils.jl")
+
+##############
+# Parameters
+##############
+case_name = ARGS[1] # case_name="pglib_opf_case300_ieee" # pglib_opf_case5_pjm
+network_formulation = ARGS[2] # network_formulation=ACPPowerModel SOCWRConicPowerModel DCPPowerModel
+icnn = parse(Bool, ARGS[3]) # icnn=true # false
+filetype = ArrowFile # ArrowFile # CSVFile
+layers = [512] # [512, 256, 64] # [256, 64, 32][1024, 1024, 1024]
+path_dataset = joinpath(dirname(@__FILE__), "data")
+case_file_path = joinpath(path_dataset, case_name)
+case_file_path_output = joinpath(case_file_path, "output", string(network_formulation))
+case_file_path_input = joinpath(case_file_path, "input", "train")
+save_file = if icnn
+    "$(case_name)_$(network_formulation)_$(replace(string(layers), ", " => "_"))_icnn"
+else
+    "$(case_name)_$(network_formulation)_$(replace(string(layers), ", " => "_"))_dnn"
+end
+
+##############
+# Load Data
+##############
+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)
+]
+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)
+]
+# batch_ids = [split(split(file, "_")[end], ".")[1] for file in file_ins]
+
+# Load input and output data tables
+if filetype === ArrowFile
+    input_table_train = Arrow.Table(file_ins)
+    output_table_train = Arrow.Table(file_outs)
+else
+    input_table_train = CSV.read(file_ins[train_idx], DataFrame)
+    output_table_train = CSV.read(file_outs[train_idx], DataFrame)
+end
+
+# Convert to dataframes
+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, :]
+
+# match
+train_table = innerjoin(input_data, output_data[!, [:id, :operational_cost]]; on=:id)
+
+input_features = names(train_table[!, Not([:id, :operational_cost])])
+
+X = Float32.(Matrix(train_table[!, input_features]))
+y = Float32.(Matrix(train_table[!, [:operational_cost]]))
+
+##############
+# Fit DNN Approximator
+##############
+
+# Define model and logger
+lg = WandbLogger(
+    project = "powermodels-obj-proxies",
+    name = "$(save_file)-$(now())",
+    config = Dict(
+        "layers" => layers,
+        "batch_size" => 32,
+        "optimiser" => "ConvexRule",
+        "learning_rate" => 0.01,
+        "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}())
+if icnn
+    optimiser=ConvexRule(
+        Flux.Optimise.Adam(get_config(lg, "learning_rate"), (0.9, 0.999), 1.0e-8, IdDict{Any,Any}())
+    )
+end
+
+# optimiser= Flux.Optimise.Adam(0.01, (0.9, 0.999), 1.0e-8, IdDict{Any,Any}())
+# nn = MultitargetNeuralNetworkRegressor(;
+#     builder=FullyConnectedBuilder(layers),
+#     rng=123,
+#     epochs=5000,
+#     optimiser=optimiser,
+#     acceleration=CUDALibs(),
+#     batch_size=32,
+#     loss=relative_rmse,
+# )
+
+nn = MultitargetNeuralNetworkRegressor(;
+    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,
+)
+
+# Constrols
+
+model_dir = joinpath(dirname(@__FILE__), "models")
+mkpath(model_dir)
+
+# save_control =
+#     MLJIteration.skip(Save(joinpath(model_dir, save_file * ".jls")), predicate=1000)
+
+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),
+    # NumberLimit(n=4),
+    # NumberSinceBest(6),
+    # PQ(; alpha=0.9, k=30),
+    # GL(; alpha=4.0),
+    InvalidValue(),
+    # Threshold(0.003),
+    TimeLimit(; t=3),
+    WithLossDo(update_loss),
+    WithReportDo(update_training_loss),
+    WithIterationsDo(update_epochs)
+]
+
+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)
+
+# Finish the run
+close(lg)
+
+# # save model if final loss is better than the best loss
+# if IterationControl.loss(iterated_pipe) < save_control.best_loss
+#     mach = mach |> cpu
+
+#     fitted_model = mach.fitresult.fitresult[1]
+
+#     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
diff --git a/examples/powermodels/visualize.jl b/examples/powermodels/visualize.jl
new file mode 100644
index 0000000..0b59871
--- /dev/null
+++ b/examples/powermodels/visualize.jl
@@ -0,0 +1,167 @@
+using Plots
+using Arrow
+using DataFrames
+using Statistics
+using LinearAlgebra
+
+cossim(x,y) = dot(x,y) / (norm(x)*norm(y))
+
+##############
+# Parameters
+##############
+network_formulation = "ACPPowerModel"
+case_name = "pglib_opf_case300_ieee"
+path_dataset = joinpath(dirname(@__FILE__), "data")
+case_file_path = joinpath(path_dataset, case_name)
+case_file_path_train = joinpath(case_file_path, "input", "train")
+case_file_path_test = joinpath(case_file_path, "input", "test")
+case_file_path_output = joinpath(case_file_path, "output", string(network_formulation))
+
+##############
+# Load Data
+##############
+iter_files_train = readdir(joinpath(case_file_path_train))
+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)
+]
+file_test = [
+    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)
+]
+
+# Load input and output data tables
+input_table_train = Arrow.Table(file_train)
+input_table_test = Arrow.Table(file_test)
+output_table = Arrow.Table(file_outs)
+
+# Convert to dataframes
+input_data_train = DataFrame(input_table_train)
+input_data_test = DataFrame(input_table_test)
+output_data = DataFrame(output_table)
+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))
+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)
+]
+output_table_soc = Arrow.Table(file_outs_soc)
+output_data_soc = DataFrame(output_table_soc)
+output_data_soc.operational_cost_soc = output_data_soc.operational_cost
+output_data_soc = output_data_soc[output_data_soc.operational_cost .> 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 = abs.(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], :]
+##############
+# Plots
+##############
+
+# Load vectors
+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
+        df[!, "load[$i]"] = zeros(size(input_data, 1))
+    end
+    for j in 1:size(input_data, 1), i in 1:num_loads
+        df[j, "load[$i]"] = sqrt(input_data[j, "pd[$i]"]^2 + input_data[j, "qd[$i]"]^2)
+    end
+    return df
+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)
+
+# Nominal Loads
+nominal_loads = Vector(load_vector_train[1, Not(:id)])
+norm_nominal_loads = norm(nominal_loads)
+
+# Load divergence
+theta_train = [acos(cossim(nominal_loads, Vector(load_vector_train[i, Not(:id)]))) for i in 2:size(load_vector_train, 1)] * 180 / pi
+norm_sim_train = [norm(Vector(load_vector_train[i, Not(:id)])) / norm_nominal_loads for i in 2:size(load_vector_train, 1)]
+
+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)]
+
+# 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)
+
+######### Plot Objective Function #########
+# Plot objective function for a single direction
+
+# Select two points in the extremes
+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
+# join input and output data
+joined_data = innerjoin(load_vector, output_data[!, [:id, :operational_cost]], on=:id)
+
+# get k extreme points
+using L2O
+using Gurobi
+function maxk(a, k)
+    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)
+    return collect(zip(b, a[b]))
+end
+k = 10
+idx_maxs = maxk(joined_data.norm_loads, k)
+idx_mins = mink(joined_data.norm_loads, k)
+
+# Objective function
+instances_in_convex_hulls = Array{DataFrame}(undef, k)
+for i in 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
+    )
+
+    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)
+end
+plt
diff --git a/examples/unitcommitment/slurm/slurm.jl b/examples/slurm/slurm.jl
similarity index 98%
rename from examples/unitcommitment/slurm/slurm.jl
rename to examples/slurm/slurm.jl
index c11ade0..3839463 100644
--- a/examples/unitcommitment/slurm/slurm.jl
+++ b/examples/slurm/slurm.jl
@@ -15,7 +15,7 @@ end
 
 using Distributed, ClusterManagers
 
-np = 20 #
+np = 50 #
 addprocs(SlurmManager(np), job_file_loc = ARGS[1]) #cpus_per_task=24, mem_per_cpu=24, partition="debug", t="08:00:00")
 
 println("We are all connected and ready.")
diff --git a/examples/training_utils.jl b/examples/training_utils.jl
new file mode 100644
index 0000000..d38e4ff
--- /dev/null
+++ b/examples/training_utils.jl
@@ -0,0 +1,81 @@
+
+
+function update_loss(loss)
+    Wandb.log(lg, Dict("metrics/loss" => loss))
+    return nothing
+end
+
+function update_training_loss(report)
+    Wandb.log(lg, Dict("metrics/training_loss" => report.training_losses[end]))
+    return nothing
+end
+
+function update_epochs(epoch)
+    Wandb.log(lg, Dict("log/epoch" => epoch)) 
+    return nothing
+end
+
+import IterationControl: needs_loss, update!, done, takedown
+
+struct WithModelLossDo{F<:Function}
+    f::F
+    stop_if_true::Bool
+    stop_message::Union{String,Nothing}
+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...)
+
+IterationControl.needs_loss(::WithModelLossDo) = true
+
+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)
+end
+
+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. " :
+            c.stop_message
+        verbosity > 0 && @info message
+        return merge(state, (log = message,))
+    else
+        return merge(state, (log = "",))
+    end
+end
+
+
+mutable struct SaveBest <: Function
+    best_loss::Float64
+    model_path::String
+    threshold::Float64
+end
+function (callback::SaveBest)(loss, ic_model)
+    update_loss(loss)
+    mach = IterationControl.expose(ic_model)
+    if loss < callback.best_loss
+        @info "best model change" callback.best_loss loss
+        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)
+    end
+    if loss < callback.threshold
+        return true
+    end
+    return false
+end
diff --git a/examples/unitcommitment/training_utils.jl b/examples/unitcommitment/training_utils.jl
deleted file mode 100644
index a69b195..0000000
--- a/examples/unitcommitment/training_utils.jl
+++ /dev/null
@@ -1,15 +0,0 @@
-
-function update_loss(loss)
-    Wandb.log(lg, Dict("metrics/loss" => loss))
-    return nothing
-end
-
-function update_training_loss(report)
-    Wandb.log(lg, Dict("metrics/training_loss" => report.training_losses[end]))
-    return nothing
-end
-
-function update_epochs(epoch)
-    Wandb.log(lg, Dict("log/epoch" => epoch)) 
-    return nothing
-end
\ No newline at end of file
diff --git a/examples/unitcommitment/uc_nn_training.jl b/examples/unitcommitment/uc_nn_training.jl
index 9c5f39c..f77ba9f 100644
--- a/examples/unitcommitment/uc_nn_training.jl
+++ b/examples/unitcommitment/uc_nn_training.jl
@@ -20,7 +20,7 @@ using Wandb, Dates, Logging
 using Statistics
 
 include(joinpath(dirname(@__FILE__), "bnb_dataset.jl"))
-include(joinpath(dirname(@__FILE__), "training_utils.jl"))
+include(joinpath(dirname(dirname(@__FILE__), "training_utils.jl")))
 
 include(joinpath(dirname(dirname(@__DIR__)), "src/cutting_planes.jl"))
 
diff --git a/src/FullyConnected.jl b/src/FullyConnected.jl
index 4730541..19b7042 100644
--- a/src/FullyConnected.jl
+++ b/src/FullyConnected.jl
@@ -68,6 +68,10 @@ end
 # Define a container to hold any optimiser specific parameters (if any):
 struct ConvexRule <: Flux.Optimise.AbstractOptimiser
     rule::Flux.Optimise.AbstractOptimiser
+    tol::Real
+end
+function ConvexRule(rule::Flux.Optimise.AbstractOptimiser; tol=1e-6)
+    return ConvexRule(rule, tol)
 end
 
 """
@@ -113,7 +117,7 @@ function MLJFlux.train!(
             return batch_loss
         end
         Flux.update!(optimiser.rule, parameters, gs)
-        make_convex!(chain)
+        make_convex!(chain; tol=optimiser.tol)
     end
     return training_loss / n_batches
 end
diff --git a/src/inconvexhull.jl b/src/inconvexhull.jl
index 1e36003..6ee4614 100644
--- a/src/inconvexhull.jl
+++ b/src/inconvexhull.jl
@@ -4,30 +4,48 @@
 
 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)
+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)
-    
-    # Create the model
-    model = JuMP.Model(solver)
+    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
+        set_attribute(model, MOI.Silent(), true)
+    end
+
     # Create the variables
-    @variable(model, lambda[1:n, 1:m] >= 0)
-    @constraint(model, convex_combination[i=1:m], sum(lambda[j, i] for j in 1:n) == 1)
+    @variable(model, lambda[1:n] >= 0)
+    @constraint(model, convex_combination, sum(lambda) == 1)
+
+    # Create POI parameters
+    @variable(model, test_set_params[1:d] in MOI.Parameter.(0.0))
     
     # slack variables
-    @variable(model, slack[1:m] >= 0)
+    @variable(model, slack[1:d])
+    @variable(model, abs_slack[1:d] >= 0)
+    @constraint(model, abs_slack .>= slack)
+    @constraint(model, abs_slack .>= -slack)
 
     # Create the constraints
-    @constraint(model, in_convex_hull[i=1:m, k=1:d], sum(lambda[j, i] * training_set[j, k] for j in 1:n) == test_set[i, k] + slack[i])
+    @constraint(model, in_convex_hull, training_set'lambda .== test_set_params .+ slack)
 
     # Create the objective
-    @objective(model, Min, sum(slack[i] for i in 1:m))
-    
-    # solve the model
-    optimize!(model)
+    @objective(model, Min, sum(abs_slack))
 
-    # return if the points are inside the convex hull
-    return isapprox.(value.(slack), 0)
+    in_convex_hull = Vector{Bool}(undef, m)
+    for i in 1:m
+        # Set the test set parameters
+        MOI.set.(model, POI.ParameterValue(), test_set_params, test_set[i, :])
+
+        # solve the model
+        optimize!(model)
+
+        # 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
diff --git a/test/inconvexhull.jl b/test/inconvexhull.jl
index 90296e1..14e6445 100644
--- a/test/inconvexhull.jl
+++ b/test/inconvexhull.jl
@@ -9,9 +9,9 @@ function test_inconvexhull()
         training_set = [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]
+        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]
+        @test inconvexhull(training_set, test_set, HiGHS.Optimizer) == [true, false, false, true, false]
     end
 end