diff --git a/dev/.documenter-siteinfo.json b/dev/.documenter-siteinfo.json
index b9c21957..6966c2d8 100644
--- a/dev/.documenter-siteinfo.json
+++ b/dev/.documenter-siteinfo.json
@@ -1 +1 @@
-{"documenter":{"julia_version":"1.11.0","generation_timestamp":"2024-10-10T16:24:19","documenter_version":"1.7.0"}}
\ No newline at end of file
+{"documenter":{"julia_version":"1.11.0","generation_timestamp":"2024-10-13T01:13:59","documenter_version":"1.7.0"}}
\ No newline at end of file
diff --git a/dev/index.html b/dev/index.html
index cae6080a..5e177064 100644
--- a/dev/index.html
+++ b/dev/index.html
@@ -1,2 +1,2 @@
 <!DOCTYPE html>
-<html lang="en"><head><meta charset="UTF-8"/><meta name="viewport" content="width=device-width, initial-scale=1.0"/><title>Home · ContinuousNormalizingFlows.jl</title><meta name="title" content="Home · ContinuousNormalizingFlows.jl"/><meta property="og:title" content="Home · ContinuousNormalizingFlows.jl"/><meta property="twitter:title" content="Home · ContinuousNormalizingFlows.jl"/><meta name="description" content="Documentation for ContinuousNormalizingFlows.jl."/><meta property="og:description" content="Documentation for ContinuousNormalizingFlows.jl."/><meta property="twitter:description" content="Documentation for ContinuousNormalizingFlows.jl."/><meta property="og:url" content="https://impICNF.github.io/ContinuousNormalizingFlows.jl/"/><meta property="twitter:url" content="https://impICNF.github.io/ContinuousNormalizingFlows.jl/"/><link rel="canonical" href="https://impICNF.github.io/ContinuousNormalizingFlows.jl/"/><script data-outdated-warner src="assets/warner.js"></script><link href="https://cdnjs.cloudflare.com/ajax/libs/lato-font/3.0.0/css/lato-font.min.css" rel="stylesheet" type="text/css"/><link href="https://cdnjs.cloudflare.com/ajax/libs/juliamono/0.050/juliamono.min.css" rel="stylesheet" type="text/css"/><link href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.4.2/css/fontawesome.min.css" rel="stylesheet" type="text/css"/><link href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.4.2/css/solid.min.css" rel="stylesheet" type="text/css"/><link href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.4.2/css/brands.min.css" rel="stylesheet" type="text/css"/><link href="https://cdnjs.cloudflare.com/ajax/libs/KaTeX/0.16.8/katex.min.css" rel="stylesheet" type="text/css"/><script>documenterBaseURL="."</script><script src="https://cdnjs.cloudflare.com/ajax/libs/require.js/2.3.6/require.min.js" data-main="assets/documenter.js"></script><script src="search_index.js"></script><script src="siteinfo.js"></script><script src="../versions.js"></script><link class="docs-theme-link" rel="stylesheet" type="text/css" href="assets/themes/catppuccin-mocha.css" data-theme-name="catppuccin-mocha"/><link class="docs-theme-link" rel="stylesheet" type="text/css" href="assets/themes/catppuccin-macchiato.css" data-theme-name="catppuccin-macchiato"/><link class="docs-theme-link" rel="stylesheet" type="text/css" href="assets/themes/catppuccin-frappe.css" data-theme-name="catppuccin-frappe"/><link class="docs-theme-link" rel="stylesheet" type="text/css" href="assets/themes/catppuccin-latte.css" data-theme-name="catppuccin-latte"/><link class="docs-theme-link" rel="stylesheet" type="text/css" href="assets/themes/documenter-dark.css" data-theme-name="documenter-dark" data-theme-primary-dark/><link class="docs-theme-link" rel="stylesheet" type="text/css" href="assets/themes/documenter-light.css" data-theme-name="documenter-light" data-theme-primary/><script src="assets/themeswap.js"></script></head><body><div id="documenter"><nav class="docs-sidebar"><div class="docs-package-name"><span class="docs-autofit"><a href>ContinuousNormalizingFlows.jl</a></span></div><button class="docs-search-query input is-rounded is-small is-clickable my-2 mx-auto py-1 px-2" id="documenter-search-query">Search docs (Ctrl + /)</button><ul class="docs-menu"><li class="is-active"><a class="tocitem" href>Home</a></li></ul><div class="docs-version-selector field has-addons"><div class="control"><span class="docs-label button is-static is-size-7">Version</span></div><div class="docs-selector control is-expanded"><div class="select is-fullwidth is-size-7"><select id="documenter-version-selector"></select></div></div></div></nav><div class="docs-main"><header class="docs-navbar"><a class="docs-sidebar-button docs-navbar-link fa-solid fa-bars is-hidden-desktop" id="documenter-sidebar-button" href="#"></a><nav class="breadcrumb"><ul class="is-hidden-mobile"><li class="is-active"><a href>Home</a></li></ul><ul class="is-hidden-tablet"><li class="is-active"><a href>Home</a></li></ul></nav><div class="docs-right"><a class="docs-navbar-link" href="https://github.com/impICNF/ContinuousNormalizingFlows.jl/blob/main/docs/src/index.md#" title="Edit source on GitHub"><span class="docs-icon fa-solid"></span></a><a class="docs-settings-button docs-navbar-link fa-solid fa-gear" id="documenter-settings-button" href="#" title="Settings"></a><a class="docs-article-toggle-button fa-solid fa-chevron-up" id="documenter-article-toggle-button" href="javascript:;" title="Collapse all docstrings"></a></div></header><article class="content" id="documenter-page"><h1 id="ContinuousNormalizingFlows"><a class="docs-heading-anchor" href="#ContinuousNormalizingFlows">ContinuousNormalizingFlows</a><a id="ContinuousNormalizingFlows-1"></a><a class="docs-heading-anchor-permalink" href="#ContinuousNormalizingFlows" title="Permalink"></a></h1><p>Documentation for <a href="https://github.com/impICNF/ContinuousNormalizingFlows.jl">ContinuousNormalizingFlows</a>.</p><ul><li><a href="#ContinuousNormalizingFlows.ContinuousNormalizingFlows"><code>ContinuousNormalizingFlows.ContinuousNormalizingFlows</code></a></li><li><a href="#ContinuousNormalizingFlows.ICNF"><code>ContinuousNormalizingFlows.ICNF</code></a></li><li><a href="#ContinuousNormalizingFlows.PlanarLayer"><code>ContinuousNormalizingFlows.PlanarLayer</code></a></li></ul><article class="docstring"><header><a class="docstring-article-toggle-button fa-solid fa-chevron-down" href="javascript:;" title="Collapse docstring"></a><a class="docstring-binding" id="ContinuousNormalizingFlows.ContinuousNormalizingFlows" href="#ContinuousNormalizingFlows.ContinuousNormalizingFlows"><code>ContinuousNormalizingFlows.ContinuousNormalizingFlows</code></a> — <span class="docstring-category">Module</span><span class="is-flex-grow-1 docstring-article-toggle-button" title="Collapse docstring"></span></header><section><div><p>Implementations of Infinitesimal Continuous Normalizing Flows Algorithms in Julia</p></div><a class="docs-sourcelink" target="_blank" href="https://github.com/impICNF/ContinuousNormalizingFlows.jl/blob/8f2e6b7c0ede18a3ff244f7e3bd4b42cb85f0c9e/src/ContinuousNormalizingFlows.jl#L74-L76">source</a></section></article><article class="docstring"><header><a class="docstring-article-toggle-button fa-solid fa-chevron-down" href="javascript:;" title="Collapse docstring"></a><a class="docstring-binding" id="ContinuousNormalizingFlows.ICNF" href="#ContinuousNormalizingFlows.ICNF"><code>ContinuousNormalizingFlows.ICNF</code></a> — <span class="docstring-category">Type</span><span class="is-flex-grow-1 docstring-article-toggle-button" title="Collapse docstring"></span></header><section><div><p>Implementation of ICNF.</p><p>Refs:</p><p><a href="https://arxiv.org/abs/1806.07366">Chen, Ricky TQ, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. &quot;Neural Ordinary Differential Equations.&quot; arXiv preprint arXiv:1806.07366 (2018).</a></p><p><a href="https://arxiv.org/abs/1810.01367">Grathwohl, Will, Ricky TQ Chen, Jesse Bettencourt, Ilya Sutskever, and David Duvenaud. &quot;Ffjord: Free-form continuous dynamics for scalable reversible generative models.&quot; arXiv preprint arXiv:1810.01367 (2018).</a></p><p><a href="https://arxiv.org/abs/2002.02798">Finlay, Chris, Jörn-Henrik Jacobsen, Levon Nurbekyan, and Adam M. Oberman. &quot;How to train your neural ODE: the world of Jacobian and kinetic regularization.&quot; arXiv preprint arXiv:2002.02798 (2020).</a></p></div><a class="docs-sourcelink" target="_blank" href="https://github.com/impICNF/ContinuousNormalizingFlows.jl/blob/8f2e6b7c0ede18a3ff244f7e3bd4b42cb85f0c9e/src/icnf.jl#L58-L68">source</a></section></article><article class="docstring"><header><a class="docstring-article-toggle-button fa-solid fa-chevron-down" href="javascript:;" title="Collapse docstring"></a><a class="docstring-binding" id="ContinuousNormalizingFlows.PlanarLayer" href="#ContinuousNormalizingFlows.PlanarLayer"><code>ContinuousNormalizingFlows.PlanarLayer</code></a> — <span class="docstring-category">Type</span><span class="is-flex-grow-1 docstring-article-toggle-button" title="Collapse docstring"></span></header><section><div><p>Implementation of Planar Layer from</p><p><a href="https://arxiv.org/abs/1806.07366">Chen, Ricky TQ, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. &quot;Neural Ordinary Differential Equations.&quot; arXiv preprint arXiv:1806.07366 (2018).</a></p></div><a class="docs-sourcelink" target="_blank" href="https://github.com/impICNF/ContinuousNormalizingFlows.jl/blob/8f2e6b7c0ede18a3ff244f7e3bd4b42cb85f0c9e/src/layers/planar_layer.jl#L1-L5">source</a></section></article></article><nav class="docs-footer"><p class="footer-message">Powered by <a href="https://github.com/JuliaDocs/Documenter.jl">Documenter.jl</a> and the <a href="https://julialang.org/">Julia Programming Language</a>.</p></nav></div><div class="modal" id="documenter-settings"><div class="modal-background"></div><div class="modal-card"><header class="modal-card-head"><p class="modal-card-title">Settings</p><button class="delete"></button></header><section class="modal-card-body"><p><label class="label">Theme</label><div class="select"><select id="documenter-themepicker"><option value="auto">Automatic (OS)</option><option value="documenter-light">documenter-light</option><option value="documenter-dark">documenter-dark</option><option value="catppuccin-latte">catppuccin-latte</option><option value="catppuccin-frappe">catppuccin-frappe</option><option value="catppuccin-macchiato">catppuccin-macchiato</option><option value="catppuccin-mocha">catppuccin-mocha</option></select></div></p><hr/><p>This document was generated with <a href="https://github.com/JuliaDocs/Documenter.jl">Documenter.jl</a> version 1.7.0 on <span class="colophon-date" title="Thursday 10 October 2024 16:24">Thursday 10 October 2024</span>. Using Julia version 1.11.0.</p></section><footer class="modal-card-foot"></footer></div></div></div></body></html>
+<html lang="en"><head><meta charset="UTF-8"/><meta name="viewport" content="width=device-width, initial-scale=1.0"/><title>Home · ContinuousNormalizingFlows.jl</title><meta name="title" content="Home · ContinuousNormalizingFlows.jl"/><meta property="og:title" content="Home · ContinuousNormalizingFlows.jl"/><meta property="twitter:title" content="Home · ContinuousNormalizingFlows.jl"/><meta name="description" content="Documentation for ContinuousNormalizingFlows.jl."/><meta property="og:description" content="Documentation for ContinuousNormalizingFlows.jl."/><meta property="twitter:description" content="Documentation for ContinuousNormalizingFlows.jl."/><meta property="og:url" content="https://impICNF.github.io/ContinuousNormalizingFlows.jl/"/><meta property="twitter:url" content="https://impICNF.github.io/ContinuousNormalizingFlows.jl/"/><link rel="canonical" href="https://impICNF.github.io/ContinuousNormalizingFlows.jl/"/><script data-outdated-warner src="assets/warner.js"></script><link href="https://cdnjs.cloudflare.com/ajax/libs/lato-font/3.0.0/css/lato-font.min.css" rel="stylesheet" type="text/css"/><link href="https://cdnjs.cloudflare.com/ajax/libs/juliamono/0.050/juliamono.min.css" rel="stylesheet" type="text/css"/><link href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.4.2/css/fontawesome.min.css" rel="stylesheet" type="text/css"/><link href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.4.2/css/solid.min.css" rel="stylesheet" type="text/css"/><link href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.4.2/css/brands.min.css" rel="stylesheet" type="text/css"/><link href="https://cdnjs.cloudflare.com/ajax/libs/KaTeX/0.16.8/katex.min.css" rel="stylesheet" type="text/css"/><script>documenterBaseURL="."</script><script src="https://cdnjs.cloudflare.com/ajax/libs/require.js/2.3.6/require.min.js" data-main="assets/documenter.js"></script><script src="search_index.js"></script><script src="siteinfo.js"></script><script src="../versions.js"></script><link class="docs-theme-link" rel="stylesheet" type="text/css" href="assets/themes/catppuccin-mocha.css" data-theme-name="catppuccin-mocha"/><link class="docs-theme-link" rel="stylesheet" type="text/css" href="assets/themes/catppuccin-macchiato.css" data-theme-name="catppuccin-macchiato"/><link class="docs-theme-link" rel="stylesheet" type="text/css" href="assets/themes/catppuccin-frappe.css" data-theme-name="catppuccin-frappe"/><link class="docs-theme-link" rel="stylesheet" type="text/css" href="assets/themes/catppuccin-latte.css" data-theme-name="catppuccin-latte"/><link class="docs-theme-link" rel="stylesheet" type="text/css" href="assets/themes/documenter-dark.css" data-theme-name="documenter-dark" data-theme-primary-dark/><link class="docs-theme-link" rel="stylesheet" type="text/css" href="assets/themes/documenter-light.css" data-theme-name="documenter-light" data-theme-primary/><script src="assets/themeswap.js"></script></head><body><div id="documenter"><nav class="docs-sidebar"><div class="docs-package-name"><span class="docs-autofit"><a href>ContinuousNormalizingFlows.jl</a></span></div><button class="docs-search-query input is-rounded is-small is-clickable my-2 mx-auto py-1 px-2" id="documenter-search-query">Search docs (Ctrl + /)</button><ul class="docs-menu"><li class="is-active"><a class="tocitem" href>Home</a></li></ul><div class="docs-version-selector field has-addons"><div class="control"><span class="docs-label button is-static is-size-7">Version</span></div><div class="docs-selector control is-expanded"><div class="select is-fullwidth is-size-7"><select id="documenter-version-selector"></select></div></div></div></nav><div class="docs-main"><header class="docs-navbar"><a class="docs-sidebar-button docs-navbar-link fa-solid fa-bars is-hidden-desktop" id="documenter-sidebar-button" href="#"></a><nav class="breadcrumb"><ul class="is-hidden-mobile"><li class="is-active"><a href>Home</a></li></ul><ul class="is-hidden-tablet"><li class="is-active"><a href>Home</a></li></ul></nav><div class="docs-right"><a class="docs-navbar-link" href="https://github.com/impICNF/ContinuousNormalizingFlows.jl/blob/main/docs/src/index.md#" title="Edit source on GitHub"><span class="docs-icon fa-solid"></span></a><a class="docs-settings-button docs-navbar-link fa-solid fa-gear" id="documenter-settings-button" href="#" title="Settings"></a><a class="docs-article-toggle-button fa-solid fa-chevron-up" id="documenter-article-toggle-button" href="javascript:;" title="Collapse all docstrings"></a></div></header><article class="content" id="documenter-page"><h1 id="ContinuousNormalizingFlows"><a class="docs-heading-anchor" href="#ContinuousNormalizingFlows">ContinuousNormalizingFlows</a><a id="ContinuousNormalizingFlows-1"></a><a class="docs-heading-anchor-permalink" href="#ContinuousNormalizingFlows" title="Permalink"></a></h1><p>Documentation for <a href="https://github.com/impICNF/ContinuousNormalizingFlows.jl">ContinuousNormalizingFlows</a>.</p><ul><li><a href="#ContinuousNormalizingFlows.ContinuousNormalizingFlows"><code>ContinuousNormalizingFlows.ContinuousNormalizingFlows</code></a></li><li><a href="#ContinuousNormalizingFlows.ICNF"><code>ContinuousNormalizingFlows.ICNF</code></a></li><li><a href="#ContinuousNormalizingFlows.PlanarLayer"><code>ContinuousNormalizingFlows.PlanarLayer</code></a></li></ul><article class="docstring"><header><a class="docstring-article-toggle-button fa-solid fa-chevron-down" href="javascript:;" title="Collapse docstring"></a><a class="docstring-binding" id="ContinuousNormalizingFlows.ContinuousNormalizingFlows" href="#ContinuousNormalizingFlows.ContinuousNormalizingFlows"><code>ContinuousNormalizingFlows.ContinuousNormalizingFlows</code></a> — <span class="docstring-category">Module</span><span class="is-flex-grow-1 docstring-article-toggle-button" title="Collapse docstring"></span></header><section><div><p>Implementations of Infinitesimal Continuous Normalizing Flows Algorithms in Julia</p></div><a class="docs-sourcelink" target="_blank" href="https://github.com/impICNF/ContinuousNormalizingFlows.jl/blob/8f2e6b7c0ede18a3ff244f7e3bd4b42cb85f0c9e/src/ContinuousNormalizingFlows.jl#L74-L76">source</a></section></article><article class="docstring"><header><a class="docstring-article-toggle-button fa-solid fa-chevron-down" href="javascript:;" title="Collapse docstring"></a><a class="docstring-binding" id="ContinuousNormalizingFlows.ICNF" href="#ContinuousNormalizingFlows.ICNF"><code>ContinuousNormalizingFlows.ICNF</code></a> — <span class="docstring-category">Type</span><span class="is-flex-grow-1 docstring-article-toggle-button" title="Collapse docstring"></span></header><section><div><p>Implementation of ICNF.</p><p>Refs:</p><p><a href="https://arxiv.org/abs/1806.07366">Chen, Ricky TQ, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. &quot;Neural Ordinary Differential Equations.&quot; arXiv preprint arXiv:1806.07366 (2018).</a></p><p><a href="https://arxiv.org/abs/1810.01367">Grathwohl, Will, Ricky TQ Chen, Jesse Bettencourt, Ilya Sutskever, and David Duvenaud. &quot;Ffjord: Free-form continuous dynamics for scalable reversible generative models.&quot; arXiv preprint arXiv:1810.01367 (2018).</a></p><p><a href="https://arxiv.org/abs/2002.02798">Finlay, Chris, Jörn-Henrik Jacobsen, Levon Nurbekyan, and Adam M. Oberman. &quot;How to train your neural ODE: the world of Jacobian and kinetic regularization.&quot; arXiv preprint arXiv:2002.02798 (2020).</a></p></div><a class="docs-sourcelink" target="_blank" href="https://github.com/impICNF/ContinuousNormalizingFlows.jl/blob/8f2e6b7c0ede18a3ff244f7e3bd4b42cb85f0c9e/src/icnf.jl#L58-L68">source</a></section></article><article class="docstring"><header><a class="docstring-article-toggle-button fa-solid fa-chevron-down" href="javascript:;" title="Collapse docstring"></a><a class="docstring-binding" id="ContinuousNormalizingFlows.PlanarLayer" href="#ContinuousNormalizingFlows.PlanarLayer"><code>ContinuousNormalizingFlows.PlanarLayer</code></a> — <span class="docstring-category">Type</span><span class="is-flex-grow-1 docstring-article-toggle-button" title="Collapse docstring"></span></header><section><div><p>Implementation of Planar Layer from</p><p><a href="https://arxiv.org/abs/1806.07366">Chen, Ricky TQ, Yulia Rubanova, Jesse Bettencourt, and David Duvenaud. &quot;Neural Ordinary Differential Equations.&quot; arXiv preprint arXiv:1806.07366 (2018).</a></p></div><a class="docs-sourcelink" target="_blank" href="https://github.com/impICNF/ContinuousNormalizingFlows.jl/blob/8f2e6b7c0ede18a3ff244f7e3bd4b42cb85f0c9e/src/layers/planar_layer.jl#L1-L5">source</a></section></article></article><nav class="docs-footer"><p class="footer-message">Powered by <a href="https://github.com/JuliaDocs/Documenter.jl">Documenter.jl</a> and the <a href="https://julialang.org/">Julia Programming Language</a>.</p></nav></div><div class="modal" id="documenter-settings"><div class="modal-background"></div><div class="modal-card"><header class="modal-card-head"><p class="modal-card-title">Settings</p><button class="delete"></button></header><section class="modal-card-body"><p><label class="label">Theme</label><div class="select"><select id="documenter-themepicker"><option value="auto">Automatic (OS)</option><option value="documenter-light">documenter-light</option><option value="documenter-dark">documenter-dark</option><option value="catppuccin-latte">catppuccin-latte</option><option value="catppuccin-frappe">catppuccin-frappe</option><option value="catppuccin-macchiato">catppuccin-macchiato</option><option value="catppuccin-mocha">catppuccin-mocha</option></select></div></p><hr/><p>This document was generated with <a href="https://github.com/JuliaDocs/Documenter.jl">Documenter.jl</a> version 1.7.0 on <span class="colophon-date" title="Sunday 13 October 2024 01:13">Sunday 13 October 2024</span>. Using Julia version 1.11.0.</p></section><footer class="modal-card-foot"></footer></div></div></div></body></html>
