diff --git a/artifacts/results/images/mnist_eccco.png b/artifacts/results/images/mnist_eccco.png
index c0a4d3cadf7419386200537f06cd135930eba6d4..d2a1a9f663808a6e1d09ce3831a95909a1b490eb 100644
Binary files a/artifacts/results/images/mnist_eccco.png and b/artifacts/results/images/mnist_eccco.png differ
diff --git a/artifacts/results/images/mnist_generated_JEM Ensemble.png b/artifacts/results/images/mnist_generated_JEM Ensemble.png
index 727cb76a00365999bddc72e8bacba2e624d2abca..cfe1c2742087a55f7affff700c0c7c29f104dde1 100644
Binary files a/artifacts/results/images/mnist_generated_JEM Ensemble.png and b/artifacts/results/images/mnist_generated_JEM Ensemble.png differ
diff --git a/artifacts/results/images/mnist_generated_JEM.png b/artifacts/results/images/mnist_generated_JEM.png
index 66e167e27972bafccda5fbca8831879ce375822c..2038ee9f5dc5873ce73fa76e0d3debfb11cfd872 100644
Binary files a/artifacts/results/images/mnist_generated_JEM.png and b/artifacts/results/images/mnist_generated_JEM.png differ
diff --git a/artifacts/results/images/mnist_generated_MLP Ensemble.png b/artifacts/results/images/mnist_generated_MLP Ensemble.png
index f9f9b1c230769ff79a383f49a5ce41d2b25556ca..5c77c2a3a5efc62d80ea29572d2e63badbe7dcf1 100644
Binary files a/artifacts/results/images/mnist_generated_MLP Ensemble.png and b/artifacts/results/images/mnist_generated_MLP Ensemble.png differ
diff --git a/artifacts/results/images/mnist_generated_MLP.png b/artifacts/results/images/mnist_generated_MLP.png
index 7ebfa028d077dd7755c0f37fb5e292f455b69582..6557ba596a962358a562d47c8137f35a9324ebeb 100644
Binary files a/artifacts/results/images/mnist_generated_MLP.png and b/artifacts/results/images/mnist_generated_MLP.png differ
diff --git a/artifacts/results/mnist_architectures.jls b/artifacts/results/mnist_architectures.jls
index 42d72be1de103a6e64215bcec3773b7664de7a96..9c2b44da5ba94dd134ff95b1c4195abffbccdcb2 100644
Binary files a/artifacts/results/mnist_architectures.jls and b/artifacts/results/mnist_architectures.jls differ
diff --git a/artifacts/results/mnist_vae.jls b/artifacts/results/mnist_vae.jls
index 12bead8bc75c668cf76bbee75526f88785b60288..82dc5667f21aa8d83d2e7b8e5718a401808bdd8e 100644
Binary files a/artifacts/results/mnist_vae.jls and b/artifacts/results/mnist_vae.jls differ
diff --git a/artifacts/results/mnist_vae_weak.jls b/artifacts/results/mnist_vae_weak.jls
index c81a673624b1d36493bfe580eda7e0ee3a03de7c..90a0c8793d037702b330e5135bb82605484964d6 100644
Binary files a/artifacts/results/mnist_vae_weak.jls and b/artifacts/results/mnist_vae_weak.jls differ
diff --git a/notebooks/mnist.qmd b/notebooks/mnist.qmd
index afe54d16141b3738fc8788e00b7c05c324883975..807d17bff860e23c9d84ff9fec5612aca3bb44bc 100644
--- a/notebooks/mnist.qmd
+++ b/notebooks/mnist.qmd
@@ -186,16 +186,16 @@ First, let's create a couple of image classifier architectures:
 ```{julia}
 # Model parameters:
 epochs = 100
-batch_size = minimum([Int(round(n_obs/10)), 128])
+batch_size = minimum([Int(round(n_obs/10)), 100])
 n_hidden = 128
 activation = Flux.swish
 builder = MLJFlux.@builder Flux.Chain(
     Dense(n_in, n_hidden, activation),
     Dense(n_hidden, n_out),
 )
-n_ens = 5                                   # number of models in ensemble
-_loss = Flux.Losses.crossentropy       # loss function
-_finaliser = Flux.softmax                        # finaliser function
+n_ens = 5                                       # number of models in ensemble
+_loss = Flux.Losses.crossentropy           # loss function
+_finaliser = Flux.softmax                            # finaliser function
 ```
 
 ```{julia}
@@ -430,7 +430,11 @@ end
 ```
 
 ```{julia}
-plt, eccco_generator, ces = _plot_eccco_mnist()
+opt = Flux.Optimise.Adam(0.01)
+Λ = [0.1,0.25,0.25]
+use_class_loss = true
+
+plt, eccco_generator, ces = _plot_eccco_mnist(opt=opt, λ=Λ, use_class_loss=use_class_loss)
 display(plt)
 savefig(plt, joinpath(output_images_path, "mnist_eccco.png"))
 ```
diff --git a/paper/paper.pdf b/paper/paper.pdf
index 5e1c3ebd24dd6a5df708ceb85fb3fa34e7587a42..4aa3ec59e609f4d6a40816456b3c6b93473d0427 100644
Binary files a/paper/paper.pdf and b/paper/paper.pdf differ