From 1353379f2e133e823c2ef0e50093528ebade9ff4 Mon Sep 17 00:00:00 2001
From: Optimox <sebastien.fischman@gmail.com>
Date: Mon, 27 Jun 2022 11:55:24 +0200
Subject: [PATCH] chore: release v4.0

---
 CHANGELOG.md                                  |    3 -
 docs/_modules/index.html                      |    8 +-
 .../pytorch_tabnet/abstract_model.html        |   83 +-
 .../pytorch_tabnet/augmentations.html         |  295 ++
 docs/_modules/pytorch_tabnet/callbacks.html   |   11 +-
 docs/_modules/pytorch_tabnet/metrics.html     |   66 +-
 .../pytorch_tabnet/multiclass_utils.html      |    7 +-
 docs/_modules/pytorch_tabnet/multitask.html   |    7 +-
 docs/_modules/pytorch_tabnet/pretraining.html |   35 +-
 .../pytorch_tabnet/pretraining_utils.html     |   11 +-
 docs/_modules/pytorch_tabnet/sparsemax.html   |    7 +-
 docs/_modules/pytorch_tabnet/tab_model.html   |    7 +-
 docs/_modules/pytorch_tabnet/tab_network.html |   43 +-
 docs/_modules/pytorch_tabnet/utils.html       |   33 +-
 docs/_modules/torch/optim/adam.html           |    7 +-
 docs/_sources/generated_docs/README.md.txt    |   97 +-
 .../generated_docs/pytorch_tabnet.rst.txt     |   57 +-
 docs/_static/pygments.css                     |   10 +-
 docs/generated_docs/README.html               |  144 +-
 docs/generated_docs/pytorch_tabnet.html       | 2432 +++++++++--------
 docs/genindex.html                            |   31 +-
 docs/index.html                               |   38 +-
 docs/py-modindex.html                         |   12 +-
 docs/search.html                              |    7 +-
 docs/searchindex.js                           |    2 +-
 pyproject.toml                                |    2 +-
 26 files changed, 2064 insertions(+), 1391 deletions(-)
 create mode 100644 docs/_modules/pytorch_tabnet/augmentations.html

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0b76baf6..4bce2dba 100755
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,7 +1,4 @@
 
-## [3.1.1](https://github.com/dreamquark-ai/tabnet/compare/v3.1.0...v3.1.1) (2021-02-02)
-
-
 ### Bug Fixes
 
 * add preds_mapper to pretraining ([76f2c85](https://github.com/dreamquark-ai/tabnet/commit/76f2c852f59c6ed2c5dc5f0766cb99310bae5f2c))
diff --git a/docs/_modules/index.html b/docs/_modules/index.html
index c2b88d03..611a62cf 100644
--- a/docs/_modules/index.html
+++ b/docs/_modules/index.html
@@ -82,14 +82,17 @@
               
             
             
-              <p class="caption"><span class="caption-text">Contents:</span></p>
+              <p><span class="caption-text">Contents:</span></p>
 <ul>
 <li class="toctree-l1"><a class="reference internal" href="../generated_docs/README.html">README</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../generated_docs/README.html#tabnet-attentive-interpretable-tabular-learning">TabNet : Attentive Interpretable Tabular Learning</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../generated_docs/README.html#installation">Installation</a></li>
-<li class="toctree-l1"><a class="reference internal" href="../generated_docs/README.html#what-problems-does-pytorch-tabnet-handles">What problems does pytorch-tabnet handles?</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../generated_docs/README.html#contributing">Contributing</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../generated_docs/README.html#what-problems-does-pytorch-tabnet-handle">What problems does pytorch-tabnet handle?</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../generated_docs/README.html#how-to-use-it">How to use it?</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../generated_docs/README.html#semi-supervised-pre-training">Semi-supervised pre-training</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../generated_docs/README.html#data-augmentation-on-the-fly">Data augmentation on the fly</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../generated_docs/README.html#easy-saving-and-loading">Easy saving and loading</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../generated_docs/README.html#useful-links">Useful links</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../generated_docs/pytorch_tabnet.html">pytorch_tabnet package</a></li>
 </ul>
@@ -155,6 +158,7 @@
             
   <h1>All modules for which code is available</h1>
 <ul><li><a href="pytorch_tabnet/abstract_model.html">pytorch_tabnet.abstract_model</a></li>
+<li><a href="pytorch_tabnet/augmentations.html">pytorch_tabnet.augmentations</a></li>
 <li><a href="pytorch_tabnet/callbacks.html">pytorch_tabnet.callbacks</a></li>
 <li><a href="pytorch_tabnet/metrics.html">pytorch_tabnet.metrics</a></li>
 <li><a href="pytorch_tabnet/multiclass_utils.html">pytorch_tabnet.multiclass_utils</a></li>
diff --git a/docs/_modules/pytorch_tabnet/abstract_model.html b/docs/_modules/pytorch_tabnet/abstract_model.html
index 21b10933..a50f2d69 100644
--- a/docs/_modules/pytorch_tabnet/abstract_model.html
+++ b/docs/_modules/pytorch_tabnet/abstract_model.html
@@ -82,14 +82,17 @@
               
             
             
-              <p class="caption"><span class="caption-text">Contents:</span></p>
+              <p><span class="caption-text">Contents:</span></p>
 <ul>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html">README</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#tabnet-attentive-interpretable-tabular-learning">TabNet : Attentive Interpretable Tabular Learning</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#installation">Installation</a></li>
-<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#what-problems-does-pytorch-tabnet-handles">What problems does pytorch-tabnet handles?</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#contributing">Contributing</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#what-problems-does-pytorch-tabnet-handle">What problems does pytorch-tabnet handle?</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#how-to-use-it">How to use it?</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#semi-supervised-pre-training">Semi-supervised pre-training</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#data-augmentation-on-the-fly">Data augmentation on the fly</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#easy-saving-and-loading">Easy saving and loading</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#useful-links">Useful links</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/pytorch_tabnet.html">pytorch_tabnet package</a></li>
 </ul>
@@ -171,6 +174,8 @@ <h1>Source code for pytorch_tabnet.abstract_model</h1><div class="highlight"><pr
     <span class="n">create_dataloaders</span><span class="p">,</span>
     <span class="n">define_device</span><span class="p">,</span>
     <span class="n">ComplexEncoder</span><span class="p">,</span>
+    <span class="n">check_input</span><span class="p">,</span>
+    <span class="n">check_warm_start</span>
 <span class="p">)</span>
 <span class="kn">from</span> <span class="nn">pytorch_tabnet.callbacks</span> <span class="kn">import</span> <span class="p">(</span>
     <span class="n">CallbackContainer</span><span class="p">,</span>
@@ -180,7 +185,7 @@ <h1>Source code for pytorch_tabnet.abstract_model</h1><div class="highlight"><pr
 <span class="p">)</span>
 <span class="kn">from</span> <span class="nn">pytorch_tabnet.metrics</span> <span class="kn">import</span> <span class="n">MetricContainer</span><span class="p">,</span> <span class="n">check_metrics</span>
 <span class="kn">from</span> <span class="nn">sklearn.base</span> <span class="kn">import</span> <span class="n">BaseEstimator</span>
-<span class="kn">from</span> <span class="nn">sklearn.utils</span> <span class="kn">import</span> <span class="n">check_array</span>
+
 <span class="kn">from</span> <span class="nn">torch.utils.data</span> <span class="kn">import</span> <span class="n">DataLoader</span>
 <span class="kn">import</span> <span class="nn">io</span>
 <span class="kn">import</span> <span class="nn">json</span>
@@ -218,6 +223,8 @@ <h1>Source code for pytorch_tabnet.abstract_model</h1><div class="highlight"><pr
     <span class="n">input_dim</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="kc">None</span>
     <span class="n">output_dim</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="kc">None</span>
     <span class="n">device_name</span><span class="p">:</span> <span class="nb">str</span> <span class="o">=</span> <span class="s2">&quot;auto&quot;</span>
+    <span class="n">n_shared_decoder</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">1</span>
+    <span class="n">n_indep_decoder</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">1</span>
 
     <span class="k">def</span> <span class="nf">__post_init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
         <span class="bp">self</span><span class="o">.</span><span class="n">batch_size</span> <span class="o">=</span> <span class="mi">1024</span>
@@ -226,7 +233,11 @@ <h1>Source code for pytorch_tabnet.abstract_model</h1><div class="highlight"><pr
         <span class="c1"># Defining device</span>
         <span class="bp">self</span><span class="o">.</span><span class="n">device</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">device</span><span class="p">(</span><span class="n">define_device</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">device_name</span><span class="p">))</span>
         <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">verbose</span> <span class="o">!=</span> <span class="mi">0</span><span class="p">:</span>
-            <span class="nb">print</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Device used : </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">device</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
+            <span class="n">warnings</span><span class="o">.</span><span class="n">warn</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Device used : </span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">device</span><span class="si">}</span><span class="s2">&quot;</span><span class="p">)</span>
+
+        <span class="c1"># create deep copies of mutable parameters</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">optimizer_fn</span> <span class="o">=</span> <span class="n">copy</span><span class="o">.</span><span class="n">deepcopy</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">optimizer_fn</span><span class="p">)</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">scheduler_fn</span> <span class="o">=</span> <span class="n">copy</span><span class="o">.</span><span class="n">deepcopy</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">scheduler_fn</span><span class="p">)</span>
 
     <span class="k">def</span> <span class="nf">__update__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="o">**</span><span class="n">kwargs</span><span class="p">):</span>
         <span class="sd">&quot;&quot;&quot;</span>
@@ -271,10 +282,12 @@ <h1>Source code for pytorch_tabnet.abstract_model</h1><div class="highlight"><pr
         <span class="n">batch_size</span><span class="o">=</span><span class="mi">1024</span><span class="p">,</span>
         <span class="n">virtual_batch_size</span><span class="o">=</span><span class="mi">128</span><span class="p">,</span>
         <span class="n">num_workers</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span>
-        <span class="n">drop_last</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
+        <span class="n">drop_last</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
         <span class="n">callbacks</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
         <span class="n">pin_memory</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
         <span class="n">from_unsupervised</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
+        <span class="n">warm_start</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
+        <span class="n">augmentations</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
     <span class="p">):</span>
         <span class="sd">&quot;&quot;&quot;Train a neural network stored in self.network</span>
 <span class="sd">        Using train_dataloader for training data and</span>
@@ -318,6 +331,8 @@ <h1>Source code for pytorch_tabnet.abstract_model</h1><div class="highlight"><pr
 <span class="sd">            Whether to set pin_memory to True or False during training</span>
 <span class="sd">        from_unsupervised: unsupervised trained model</span>
 <span class="sd">            Use a previously self supervised model as starting weights</span>
+<span class="sd">        warm_start: bool</span>
+<span class="sd">            If True, current model parameters are used to start training</span>
 <span class="sd">        &quot;&quot;&quot;</span>
         <span class="c1"># update model name</span>
 
@@ -330,6 +345,11 @@ <h1>Source code for pytorch_tabnet.abstract_model</h1><div class="highlight"><pr
         <span class="bp">self</span><span class="o">.</span><span class="n">input_dim</span> <span class="o">=</span> <span class="n">X_train</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span>
         <span class="bp">self</span><span class="o">.</span><span class="n">_stop_training</span> <span class="o">=</span> <span class="kc">False</span>
         <span class="bp">self</span><span class="o">.</span><span class="n">pin_memory</span> <span class="o">=</span> <span class="n">pin_memory</span> <span class="ow">and</span> <span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">device</span><span class="o">.</span><span class="n">type</span> <span class="o">!=</span> <span class="s2">&quot;cpu&quot;</span><span class="p">)</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">augmentations</span> <span class="o">=</span> <span class="n">augmentations</span>
+
+        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">augmentations</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
+            <span class="c1"># This ensure reproducibility</span>
+            <span class="bp">self</span><span class="o">.</span><span class="n">augmentations</span><span class="o">.</span><span class="n">_set_seed</span><span class="p">()</span>
 
         <span class="n">eval_set</span> <span class="o">=</span> <span class="n">eval_set</span> <span class="k">if</span> <span class="n">eval_set</span> <span class="k">else</span> <span class="p">[]</span>
 
@@ -338,7 +358,8 @@ <h1>Source code for pytorch_tabnet.abstract_model</h1><div class="highlight"><pr
         <span class="k">else</span><span class="p">:</span>
             <span class="bp">self</span><span class="o">.</span><span class="n">loss_fn</span> <span class="o">=</span> <span class="n">loss_fn</span>
 
-        <span class="n">check_array</span><span class="p">(</span><span class="n">X_train</span><span class="p">)</span>
+        <span class="n">check_input</span><span class="p">(</span><span class="n">X_train</span><span class="p">)</span>
+        <span class="n">check_warm_start</span><span class="p">(</span><span class="n">warm_start</span><span class="p">,</span> <span class="n">from_unsupervised</span><span class="p">)</span>
 
         <span class="bp">self</span><span class="o">.</span><span class="n">update_fit_params</span><span class="p">(</span>
             <span class="n">X_train</span><span class="p">,</span>
@@ -358,7 +379,8 @@ <h1>Source code for pytorch_tabnet.abstract_model</h1><div class="highlight"><pr
             <span class="c1"># Update parameters to match self pretraining</span>
             <span class="bp">self</span><span class="o">.</span><span class="n">__update__</span><span class="p">(</span><span class="o">**</span><span class="n">from_unsupervised</span><span class="o">.</span><span class="n">get_params</span><span class="p">())</span>
 
-        <span class="k">if</span> <span class="ow">not</span> <span class="nb">hasattr</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="s2">&quot;network&quot;</span><span class="p">):</span>
+        <span class="k">if</span> <span class="ow">not</span> <span class="nb">hasattr</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="s2">&quot;network&quot;</span><span class="p">)</span> <span class="ow">or</span> <span class="ow">not</span> <span class="n">warm_start</span><span class="p">:</span>
+            <span class="c1"># model has never been fitted before of warm_start is False</span>
             <span class="bp">self</span><span class="o">.</span><span class="n">_set_network</span><span class="p">()</span>
         <span class="bp">self</span><span class="o">.</span><span class="n">_update_network_params</span><span class="p">()</span>
         <span class="bp">self</span><span class="o">.</span><span class="n">_set_metrics</span><span class="p">(</span><span class="n">eval_metric</span><span class="p">,</span> <span class="n">eval_names</span><span class="p">)</span>
@@ -366,9 +388,8 @@ <h1>Source code for pytorch_tabnet.abstract_model</h1><div class="highlight"><pr
         <span class="bp">self</span><span class="o">.</span><span class="n">_set_callbacks</span><span class="p">(</span><span class="n">callbacks</span><span class="p">)</span>
 
         <span class="k">if</span> <span class="n">from_unsupervised</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
-            <span class="nb">print</span><span class="p">(</span><span class="s2">&quot;Loading weights from unsupervised pretraining&quot;</span><span class="p">)</span>
             <span class="bp">self</span><span class="o">.</span><span class="n">load_weights_from_unsupervised</span><span class="p">(</span><span class="n">from_unsupervised</span><span class="p">)</span>
-
+            <span class="n">warnings</span><span class="o">.</span><span class="n">warn</span><span class="p">(</span><span class="s2">&quot;Loading weights from unsupervised pretraining&quot;</span><span class="p">)</span>
         <span class="c1"># Call method on_train_begin for all callbacks</span>
         <span class="bp">self</span><span class="o">.</span><span class="n">_callback_container</span><span class="o">.</span><span class="n">on_train_begin</span><span class="p">()</span>
 
@@ -397,7 +418,7 @@ <h1>Source code for pytorch_tabnet.abstract_model</h1><div class="highlight"><pr
         <span class="bp">self</span><span class="o">.</span><span class="n">network</span><span class="o">.</span><span class="n">eval</span><span class="p">()</span>
 
         <span class="c1"># compute feature importance once the best model is defined</span>
-        <span class="bp">self</span><span class="o">.</span><span class="n">_compute_feature_importances</span><span class="p">(</span><span class="n">train_dataloader</span><span class="p">)</span></div>
+        <span class="bp">self</span><span class="o">.</span><span class="n">feature_importances_</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_compute_feature_importances</span><span class="p">(</span><span class="n">X_train</span><span class="p">)</span></div>
 
 <div class="viewcode-block" id="TabModel.predict"><a class="viewcode-back" href="../../generated_docs/pytorch_tabnet.html#pytorch_tabnet.abstract_model.TabModel.predict">[docs]</a>    <span class="k">def</span> <span class="nf">predict</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">X</span><span class="p">):</span>
         <span class="sd">&quot;&quot;&quot;</span>
@@ -429,7 +450,7 @@ <h1>Source code for pytorch_tabnet.abstract_model</h1><div class="highlight"><pr
         <span class="n">res</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">vstack</span><span class="p">(</span><span class="n">results</span><span class="p">)</span>
         <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">predict_func</span><span class="p">(</span><span class="n">res</span><span class="p">)</span></div>
 
-<div class="viewcode-block" id="TabModel.explain"><a class="viewcode-back" href="../../generated_docs/pytorch_tabnet.html#pytorch_tabnet.abstract_model.TabModel.explain">[docs]</a>    <span class="k">def</span> <span class="nf">explain</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">X</span><span class="p">):</span>
+<div class="viewcode-block" id="TabModel.explain"><a class="viewcode-back" href="../../generated_docs/pytorch_tabnet.html#pytorch_tabnet.abstract_model.TabModel.explain">[docs]</a>    <span class="k">def</span> <span class="nf">explain</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">X</span><span class="p">,</span> <span class="n">normalize</span><span class="o">=</span><span class="kc">False</span><span class="p">):</span>
         <span class="sd">&quot;&quot;&quot;</span>
 <span class="sd">        Return local explanation</span>
 
@@ -437,6 +458,8 @@ <h1>Source code for pytorch_tabnet.abstract_model</h1><div class="highlight"><pr
 <span class="sd">        ----------</span>
 <span class="sd">        X : tensor: `torch.Tensor`</span>
 <span class="sd">            Input data</span>
+<span class="sd">        normalize : bool (default False)</span>
+<span class="sd">            Wheter to normalize so that sum of features are equal to 1</span>
 
 <span class="sd">        Returns</span>
 <span class="sd">        -------</span>
@@ -464,9 +487,9 @@ <h1>Source code for pytorch_tabnet.abstract_model</h1><div class="highlight"><pr
                     <span class="n">value</span><span class="o">.</span><span class="n">cpu</span><span class="p">()</span><span class="o">.</span><span class="n">detach</span><span class="p">()</span><span class="o">.</span><span class="n">numpy</span><span class="p">(),</span> <span class="bp">self</span><span class="o">.</span><span class="n">reducing_matrix</span>
                 <span class="p">)</span>
 
-            <span class="n">res_explain</span><span class="o">.</span><span class="n">append</span><span class="p">(</span>
-                <span class="n">csc_matrix</span><span class="o">.</span><span class="n">dot</span><span class="p">(</span><span class="n">M_explain</span><span class="o">.</span><span class="n">cpu</span><span class="p">()</span><span class="o">.</span><span class="n">detach</span><span class="p">()</span><span class="o">.</span><span class="n">numpy</span><span class="p">(),</span> <span class="bp">self</span><span class="o">.</span><span class="n">reducing_matrix</span><span class="p">)</span>
-            <span class="p">)</span>
+            <span class="n">original_feat_explain</span> <span class="o">=</span> <span class="n">csc_matrix</span><span class="o">.</span><span class="n">dot</span><span class="p">(</span><span class="n">M_explain</span><span class="o">.</span><span class="n">cpu</span><span class="p">()</span><span class="o">.</span><span class="n">detach</span><span class="p">()</span><span class="o">.</span><span class="n">numpy</span><span class="p">(),</span>
+                                                   <span class="bp">self</span><span class="o">.</span><span class="n">reducing_matrix</span><span class="p">)</span>
+            <span class="n">res_explain</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">original_feat_explain</span><span class="p">)</span>
 
             <span class="k">if</span> <span class="n">batch_nb</span> <span class="o">==</span> <span class="mi">0</span><span class="p">:</span>
                 <span class="n">res_masks</span> <span class="o">=</span> <span class="n">masks</span>
@@ -476,6 +499,9 @@ <h1>Source code for pytorch_tabnet.abstract_model</h1><div class="highlight"><pr
 
         <span class="n">res_explain</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">vstack</span><span class="p">(</span><span class="n">res_explain</span><span class="p">)</span>
 
+        <span class="k">if</span> <span class="n">normalize</span><span class="p">:</span>
+            <span class="n">res_explain</span> <span class="o">/=</span> <span class="n">np</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">res_explain</span><span class="p">,</span> <span class="n">axis</span><span class="o">=</span><span class="mi">1</span><span class="p">)[:,</span> <span class="kc">None</span><span class="p">]</span>
+
         <span class="k">return</span> <span class="n">res_explain</span><span class="p">,</span> <span class="n">res_masks</span></div>
 
 <div class="viewcode-block" id="TabModel.load_weights_from_unsupervised"><a class="viewcode-back" href="../../generated_docs/pytorch_tabnet.html#pytorch_tabnet.abstract_model.TabModel.load_weights_from_unsupervised">[docs]</a>    <span class="k">def</span> <span class="nf">load_weights_from_unsupervised</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">unsupervised_model</span><span class="p">):</span>
@@ -621,6 +647,9 @@ <h1>Source code for pytorch_tabnet.abstract_model</h1><div class="highlight"><pr
         <span class="n">X</span> <span class="o">=</span> <span class="n">X</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">device</span><span class="p">)</span><span class="o">.</span><span class="n">float</span><span class="p">()</span>
         <span class="n">y</span> <span class="o">=</span> <span class="n">y</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">device</span><span class="p">)</span><span class="o">.</span><span class="n">float</span><span class="p">()</span>
 
+        <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">augmentations</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
+            <span class="n">X</span><span class="p">,</span> <span class="n">y</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">augmentations</span><span class="p">(</span><span class="n">X</span><span class="p">,</span> <span class="n">y</span><span class="p">)</span>
+
         <span class="k">for</span> <span class="n">param</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">network</span><span class="o">.</span><span class="n">parameters</span><span class="p">():</span>
             <span class="n">param</span><span class="o">.</span><span class="n">grad</span> <span class="o">=</span> <span class="kc">None</span>
 
@@ -628,7 +657,7 @@ <h1>Source code for pytorch_tabnet.abstract_model</h1><div class="highlight"><pr
 
         <span class="n">loss</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">compute_loss</span><span class="p">(</span><span class="n">output</span><span class="p">,</span> <span class="n">y</span><span class="p">)</span>
         <span class="c1"># Add the overall sparsity loss</span>
-        <span class="n">loss</span> <span class="o">-=</span> <span class="bp">self</span><span class="o">.</span><span class="n">lambda_sparse</span> <span class="o">*</span> <span class="n">M_loss</span>
+        <span class="n">loss</span> <span class="o">=</span> <span class="n">loss</span> <span class="o">-</span> <span class="bp">self</span><span class="o">.</span><span class="n">lambda_sparse</span> <span class="o">*</span> <span class="n">M_loss</span>
 
         <span class="c1"># Perform backward pass and optimization</span>
         <span class="n">loss</span><span class="o">.</span><span class="n">backward</span><span class="p">()</span>
@@ -698,6 +727,7 @@ <h1>Source code for pytorch_tabnet.abstract_model</h1><div class="highlight"><pr
 
     <span class="k">def</span> <span class="nf">_set_network</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
         <span class="sd">&quot;&quot;&quot;Setup the network and explain matrix.&quot;&quot;&quot;</span>
+        <span class="n">torch</span><span class="o">.</span><span class="n">manual_seed</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">seed</span><span class="p">)</span>
         <span class="bp">self</span><span class="o">.</span><span class="n">network</span> <span class="o">=</span> <span class="n">tab_network</span><span class="o">.</span><span class="n">TabNet</span><span class="p">(</span>
             <span class="bp">self</span><span class="o">.</span><span class="n">input_dim</span><span class="p">,</span>
             <span class="bp">self</span><span class="o">.</span><span class="n">output_dim</span><span class="p">,</span>
@@ -778,9 +808,9 @@ <h1>Source code for pytorch_tabnet.abstract_model</h1><div class="highlight"><pr
             <span class="p">)</span>
             <span class="n">callbacks</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">early_stopping</span><span class="p">)</span>
         <span class="k">else</span><span class="p">:</span>
-            <span class="nb">print</span><span class="p">(</span>
-                <span class="s2">&quot;No early stopping will be performed, last training weights will be used.&quot;</span>
-            <span class="p">)</span>
+            <span class="n">wrn_msg</span> <span class="o">=</span> <span class="s2">&quot;No early stopping will be performed, last training weights will be used.&quot;</span>
+            <span class="n">warnings</span><span class="o">.</span><span class="n">warn</span><span class="p">(</span><span class="n">wrn_msg</span><span class="p">)</span>
+
         <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">scheduler_fn</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
             <span class="c1"># Add LR Scheduler call_back</span>
             <span class="n">is_batch_level</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">scheduler_params</span><span class="o">.</span><span class="n">pop</span><span class="p">(</span><span class="s2">&quot;is_batch_level&quot;</span><span class="p">,</span> <span class="kc">False</span><span class="p">)</span>
@@ -842,7 +872,7 @@ <h1>Source code for pytorch_tabnet.abstract_model</h1><div class="highlight"><pr
         <span class="p">)</span>
         <span class="k">return</span> <span class="n">train_dataloader</span><span class="p">,</span> <span class="n">valid_dataloaders</span>
 
-    <span class="k">def</span> <span class="nf">_compute_feature_importances</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">loader</span><span class="p">):</span>
+    <span class="k">def</span> <span class="nf">_compute_feature_importances</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">X</span><span class="p">):</span>
         <span class="sd">&quot;&quot;&quot;Compute global feature importance.</span>
 
 <span class="sd">        Parameters</span>
@@ -851,17 +881,10 @@ <h1>Source code for pytorch_tabnet.abstract_model</h1><div class="highlight"><pr
 <span class="sd">            Pytorch dataloader.</span>
 
 <span class="sd">        &quot;&quot;&quot;</span>
-        <span class="bp">self</span><span class="o">.</span><span class="n">network</span><span class="o">.</span><span class="n">eval</span><span class="p">()</span>
-        <span class="n">feature_importances_</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">zeros</span><span class="p">((</span><span class="bp">self</span><span class="o">.</span><span class="n">network</span><span class="o">.</span><span class="n">post_embed_dim</span><span class="p">))</span>
-        <span class="k">for</span> <span class="n">data</span><span class="p">,</span> <span class="n">targets</span> <span class="ow">in</span> <span class="n">loader</span><span class="p">:</span>
-            <span class="n">data</span> <span class="o">=</span> <span class="n">data</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">device</span><span class="p">)</span><span class="o">.</span><span class="n">float</span><span class="p">()</span>
-            <span class="n">M_explain</span><span class="p">,</span> <span class="n">masks</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">network</span><span class="o">.</span><span class="n">forward_masks</span><span class="p">(</span><span class="n">data</span><span class="p">)</span>
-            <span class="n">feature_importances_</span> <span class="o">+=</span> <span class="n">M_explain</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">dim</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span><span class="o">.</span><span class="n">cpu</span><span class="p">()</span><span class="o">.</span><span class="n">detach</span><span class="p">()</span><span class="o">.</span><span class="n">numpy</span><span class="p">()</span>
-
-        <span class="n">feature_importances_</span> <span class="o">=</span> <span class="n">csc_matrix</span><span class="o">.</span><span class="n">dot</span><span class="p">(</span>
-            <span class="n">feature_importances_</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">reducing_matrix</span>
-        <span class="p">)</span>
-        <span class="bp">self</span><span class="o">.</span><span class="n">feature_importances_</span> <span class="o">=</span> <span class="n">feature_importances_</span> <span class="o">/</span> <span class="n">np</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">feature_importances_</span><span class="p">)</span>
+        <span class="n">M_explain</span><span class="p">,</span> <span class="n">_</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">explain</span><span class="p">(</span><span class="n">X</span><span class="p">,</span> <span class="n">normalize</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
+        <span class="n">sum_explain</span> <span class="o">=</span> <span class="n">M_explain</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span>
+        <span class="n">feature_importances_</span> <span class="o">=</span> <span class="n">sum_explain</span> <span class="o">/</span> <span class="n">np</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">sum_explain</span><span class="p">)</span>
+        <span class="k">return</span> <span class="n">feature_importances_</span>
 
     <span class="k">def</span> <span class="nf">_update_network_params</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
         <span class="bp">self</span><span class="o">.</span><span class="n">network</span><span class="o">.</span><span class="n">virtual_batch_size</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">virtual_batch_size</span>
diff --git a/docs/_modules/pytorch_tabnet/augmentations.html b/docs/_modules/pytorch_tabnet/augmentations.html
new file mode 100644
index 00000000..48bc1f47
--- /dev/null
+++ b/docs/_modules/pytorch_tabnet/augmentations.html
@@ -0,0 +1,295 @@
+
+
+<!DOCTYPE html>
+<html class="writer-html5" lang="en" >
+<head>
+  <meta charset="utf-8">
+  
+  <meta name="viewport" content="width=device-width, initial-scale=1.0">
+  
+  <title>pytorch_tabnet.augmentations &mdash; pytorch_tabnet  documentation</title>
+  
+
+  
+  <link rel="stylesheet" href="../../_static/css/theme.css" type="text/css" />
+  <link rel="stylesheet" href="../../_static/pygments.css" type="text/css" />
+  <link rel="stylesheet" href="../../_static/graphviz.css" type="text/css" />
+  <link rel="stylesheet" href="../../_static/./default.css" type="text/css" />
+
+  
+  
+  
+  
+
+  
+  <!--[if lt IE 9]>
+    <script src="../../_static/js/html5shiv.min.js"></script>
+  <![endif]-->
+  
+    
+      <script type="text/javascript" id="documentation_options" data-url_root="../../" src="../../_static/documentation_options.js"></script>
+        <script src="../../_static/jquery.js"></script>
+        <script src="../../_static/underscore.js"></script>
+        <script src="../../_static/doctools.js"></script>
+        <script src="../../_static/language_data.js"></script>
+    
+    <script type="text/javascript" src="../../_static/js/theme.js"></script>
+
+    
+    <link rel="index" title="Index" href="../../genindex.html" />
+    <link rel="search" title="Search" href="../../search.html" /> 
+</head>
+
+<body class="wy-body-for-nav">
+
+   
+  <div class="wy-grid-for-nav">
+    
+    <nav data-toggle="wy-nav-shift" class="wy-nav-side">
+      <div class="wy-side-scroll">
+        <div class="wy-side-nav-search" >
+          
+
+          
+            <a href="../../index.html" class="icon icon-home" alt="Documentation Home"> pytorch_tabnet
+          
+
+          
+          </a>
+
+          
+            
+            
+          
+
+          
+<div role="search">
+  <form id="rtd-search-form" class="wy-form" action="../../search.html" method="get">
+    <input type="text" name="q" placeholder="Search docs" />
+    <input type="hidden" name="check_keywords" value="yes" />
+    <input type="hidden" name="area" value="default" />
+  </form>
+</div>
+
+          
+        </div>
+
+        
+        <div class="wy-menu wy-menu-vertical" data-spy="affix" role="navigation" aria-label="main navigation">
+          
+            
+            
+              
+            
+            
+              <p><span class="caption-text">Contents:</span></p>
+<ul>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html">README</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#tabnet-attentive-interpretable-tabular-learning">TabNet : Attentive Interpretable Tabular Learning</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#installation">Installation</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#contributing">Contributing</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#what-problems-does-pytorch-tabnet-handle">What problems does pytorch-tabnet handle?</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#how-to-use-it">How to use it?</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#semi-supervised-pre-training">Semi-supervised pre-training</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#data-augmentation-on-the-fly">Data augmentation on the fly</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#easy-saving-and-loading">Easy saving and loading</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#useful-links">Useful links</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/pytorch_tabnet.html">pytorch_tabnet package</a></li>
+</ul>
+
+            
+          
+        </div>
+        
+      </div>
+    </nav>
+
+    <section data-toggle="wy-nav-shift" class="wy-nav-content-wrap">
+
+      
+      <nav class="wy-nav-top" aria-label="top navigation">
+        
+          <i data-toggle="wy-nav-top" class="fa fa-bars"></i>
+          <a href="../../index.html">pytorch_tabnet</a>
+        
+      </nav>
+
+
+      <div class="wy-nav-content">
+        
+        <div class="rst-content">
+        
+          
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+<div role="navigation" aria-label="breadcrumbs navigation">
+
+  <ul class="wy-breadcrumbs">
+    
+      <li><a href="../../index.html" class="icon icon-home"></a> &raquo;</li>
+        
+          <li><a href="../index.html">Module code</a> &raquo;</li>
+        
+      <li>pytorch_tabnet.augmentations</li>
+    
+    
+      <li class="wy-breadcrumbs-aside">
+        
+      </li>
+    
+  </ul>
+
+  
+  <hr/>
+</div>
+          <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
+           <div itemprop="articleBody">
+            
+  <h1>Source code for pytorch_tabnet.augmentations</h1><div class="highlight"><pre>
+<span></span><span class="kn">import</span> <span class="nn">torch</span>
+<span class="kn">from</span> <span class="nn">pytorch_tabnet.utils</span> <span class="kn">import</span> <span class="n">define_device</span>
+<span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
+
+
+<div class="viewcode-block" id="RegressionSMOTE"><a class="viewcode-back" href="../../generated_docs/pytorch_tabnet.html#pytorch_tabnet.augmentations.RegressionSMOTE">[docs]</a><span class="k">class</span> <span class="nc">RegressionSMOTE</span><span class="p">():</span>
+    <span class="sd">&quot;&quot;&quot;</span>
+<span class="sd">    Apply SMOTE</span>
+
+<span class="sd">    This will average a percentage p of the elements in the batch with other elements.</span>
+<span class="sd">    The target will be averaged as well (this might work with binary classification</span>
+<span class="sd">    and certain loss), following a beta distribution.</span>
+<span class="sd">    &quot;&quot;&quot;</span>
+    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">device_name</span><span class="o">=</span><span class="s2">&quot;auto&quot;</span><span class="p">,</span> <span class="n">p</span><span class="o">=</span><span class="mf">0.8</span><span class="p">,</span> <span class="n">alpha</span><span class="o">=</span><span class="mf">0.5</span><span class="p">,</span> <span class="n">beta</span><span class="o">=</span><span class="mf">0.5</span><span class="p">,</span> <span class="n">seed</span><span class="o">=</span><span class="mi">0</span><span class="p">):</span>
+        <span class="s2">&quot;&quot;</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">seed</span> <span class="o">=</span> <span class="n">seed</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">_set_seed</span><span class="p">()</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">device</span> <span class="o">=</span> <span class="n">define_device</span><span class="p">(</span><span class="n">device_name</span><span class="p">)</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">alpha</span> <span class="o">=</span> <span class="n">alpha</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">beta</span> <span class="o">=</span> <span class="n">beta</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">p</span> <span class="o">=</span> <span class="n">p</span>
+        <span class="k">if</span> <span class="p">(</span><span class="n">p</span> <span class="o">&lt;</span> <span class="mf">0.</span><span class="p">)</span> <span class="ow">or</span> <span class="p">(</span><span class="n">p</span> <span class="o">&gt;</span> <span class="mf">1.0</span><span class="p">):</span>
+            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Value of p should be between 0. and 1.&quot;</span><span class="p">)</span>
+
+    <span class="k">def</span> <span class="nf">_set_seed</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
+        <span class="n">torch</span><span class="o">.</span><span class="n">manual_seed</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">seed</span><span class="p">)</span>
+        <span class="n">np</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">seed</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">seed</span><span class="p">)</span>
+        <span class="k">return</span>
+
+    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">X</span><span class="p">,</span> <span class="n">y</span><span class="p">):</span>
+        <span class="n">batch_size</span> <span class="o">=</span> <span class="n">X</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
+        <span class="n">random_values</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">rand</span><span class="p">(</span><span class="n">batch_size</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">device</span><span class="p">)</span>
+        <span class="n">idx_to_change</span> <span class="o">=</span> <span class="n">random_values</span> <span class="o">&lt;</span> <span class="bp">self</span><span class="o">.</span><span class="n">p</span>
+
+        <span class="c1"># ensure that first element to switch has probability &gt; 0.5</span>
+        <span class="n">np_betas</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">beta</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">alpha</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">beta</span><span class="p">,</span> <span class="n">batch_size</span><span class="p">)</span> <span class="o">/</span> <span class="mi">2</span> <span class="o">+</span> <span class="mf">0.5</span>
+        <span class="n">random_betas</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">from_numpy</span><span class="p">(</span><span class="n">np_betas</span><span class="p">)</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">device</span><span class="p">)</span><span class="o">.</span><span class="n">float</span><span class="p">()</span>
+        <span class="n">index_permute</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">randperm</span><span class="p">(</span><span class="n">batch_size</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">device</span><span class="p">)</span>
+
+        <span class="n">X</span><span class="p">[</span><span class="n">idx_to_change</span><span class="p">]</span> <span class="o">=</span> <span class="n">random_betas</span><span class="p">[</span><span class="n">idx_to_change</span><span class="p">,</span> <span class="kc">None</span><span class="p">]</span> <span class="o">*</span> <span class="n">X</span><span class="p">[</span><span class="n">idx_to_change</span><span class="p">]</span>
+        <span class="n">X</span><span class="p">[</span><span class="n">idx_to_change</span><span class="p">]</span> <span class="o">+=</span> <span class="p">(</span><span class="mi">1</span> <span class="o">-</span> <span class="n">random_betas</span><span class="p">[</span><span class="n">idx_to_change</span><span class="p">,</span> <span class="kc">None</span><span class="p">])</span> <span class="o">*</span> <span class="n">X</span><span class="p">[</span><span class="n">index_permute</span><span class="p">][</span><span class="n">idx_to_change</span><span class="p">]</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="n">X</span><span class="p">[</span><span class="n">idx_to_change</span><span class="p">]</span><span class="o">.</span><span class="n">size</span><span class="p">())</span> <span class="c1"># noqa</span>
+
+        <span class="n">y</span><span class="p">[</span><span class="n">idx_to_change</span><span class="p">]</span> <span class="o">=</span> <span class="n">random_betas</span><span class="p">[</span><span class="n">idx_to_change</span><span class="p">,</span> <span class="kc">None</span><span class="p">]</span> <span class="o">*</span> <span class="n">y</span><span class="p">[</span><span class="n">idx_to_change</span><span class="p">]</span>
+        <span class="n">y</span><span class="p">[</span><span class="n">idx_to_change</span><span class="p">]</span> <span class="o">+=</span> <span class="p">(</span><span class="mi">1</span> <span class="o">-</span> <span class="n">random_betas</span><span class="p">[</span><span class="n">idx_to_change</span><span class="p">,</span> <span class="kc">None</span><span class="p">])</span> <span class="o">*</span> <span class="n">y</span><span class="p">[</span><span class="n">index_permute</span><span class="p">][</span><span class="n">idx_to_change</span><span class="p">]</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="n">y</span><span class="p">[</span><span class="n">idx_to_change</span><span class="p">]</span><span class="o">.</span><span class="n">size</span><span class="p">())</span> <span class="c1"># noqa</span>
+
+        <span class="k">return</span> <span class="n">X</span><span class="p">,</span> <span class="n">y</span></div>
+
+
+<div class="viewcode-block" id="ClassificationSMOTE"><a class="viewcode-back" href="../../generated_docs/pytorch_tabnet.html#pytorch_tabnet.augmentations.ClassificationSMOTE">[docs]</a><span class="k">class</span> <span class="nc">ClassificationSMOTE</span><span class="p">():</span>
+    <span class="sd">&quot;&quot;&quot;</span>
+<span class="sd">    Apply SMOTE for classification tasks.</span>
+
+<span class="sd">    This will average a percentage p of the elements in the batch with other elements.</span>
+<span class="sd">    The target will stay unchanged and keep the value of the most important row in the mix.</span>
+<span class="sd">    &quot;&quot;&quot;</span>
+    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">device_name</span><span class="o">=</span><span class="s2">&quot;auto&quot;</span><span class="p">,</span> <span class="n">p</span><span class="o">=</span><span class="mf">0.8</span><span class="p">,</span> <span class="n">alpha</span><span class="o">=</span><span class="mf">0.5</span><span class="p">,</span> <span class="n">beta</span><span class="o">=</span><span class="mf">0.5</span><span class="p">,</span> <span class="n">seed</span><span class="o">=</span><span class="mi">0</span><span class="p">):</span>
+        <span class="s2">&quot;&quot;</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">seed</span> <span class="o">=</span> <span class="n">seed</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">_set_seed</span><span class="p">()</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">device</span> <span class="o">=</span> <span class="n">define_device</span><span class="p">(</span><span class="n">device_name</span><span class="p">)</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">alpha</span> <span class="o">=</span> <span class="n">alpha</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">beta</span> <span class="o">=</span> <span class="n">beta</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">p</span> <span class="o">=</span> <span class="n">p</span>
+        <span class="k">if</span> <span class="p">(</span><span class="n">p</span> <span class="o">&lt;</span> <span class="mf">0.</span><span class="p">)</span> <span class="ow">or</span> <span class="p">(</span><span class="n">p</span> <span class="o">&gt;</span> <span class="mf">1.0</span><span class="p">):</span>
+            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;Value of p should be between 0. and 1.&quot;</span><span class="p">)</span>
+
+    <span class="k">def</span> <span class="nf">_set_seed</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
+        <span class="n">torch</span><span class="o">.</span><span class="n">manual_seed</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">seed</span><span class="p">)</span>
+        <span class="n">np</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">seed</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">seed</span><span class="p">)</span>
+        <span class="k">return</span>
+
+    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">X</span><span class="p">,</span> <span class="n">y</span><span class="p">):</span>
+        <span class="n">batch_size</span> <span class="o">=</span> <span class="n">X</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span>
+        <span class="n">random_values</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">rand</span><span class="p">(</span><span class="n">batch_size</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">device</span><span class="p">)</span>
+        <span class="n">idx_to_change</span> <span class="o">=</span> <span class="n">random_values</span> <span class="o">&lt;</span> <span class="bp">self</span><span class="o">.</span><span class="n">p</span>
+
+        <span class="c1"># ensure that first element to switch has probability &gt; 0.5</span>
+        <span class="n">np_betas</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">random</span><span class="o">.</span><span class="n">beta</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">alpha</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">beta</span><span class="p">,</span> <span class="n">batch_size</span><span class="p">)</span> <span class="o">/</span> <span class="mi">2</span> <span class="o">+</span> <span class="mf">0.5</span>
+        <span class="n">random_betas</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">from_numpy</span><span class="p">(</span><span class="n">np_betas</span><span class="p">)</span><span class="o">.</span><span class="n">to</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">device</span><span class="p">)</span><span class="o">.</span><span class="n">float</span><span class="p">()</span>
+        <span class="n">index_permute</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">randperm</span><span class="p">(</span><span class="n">batch_size</span><span class="p">,</span> <span class="n">device</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">device</span><span class="p">)</span>
+
+        <span class="n">X</span><span class="p">[</span><span class="n">idx_to_change</span><span class="p">]</span> <span class="o">=</span> <span class="n">random_betas</span><span class="p">[</span><span class="n">idx_to_change</span><span class="p">,</span> <span class="kc">None</span><span class="p">]</span> <span class="o">*</span> <span class="n">X</span><span class="p">[</span><span class="n">idx_to_change</span><span class="p">]</span>
+        <span class="n">X</span><span class="p">[</span><span class="n">idx_to_change</span><span class="p">]</span> <span class="o">+=</span> <span class="p">(</span><span class="mi">1</span> <span class="o">-</span> <span class="n">random_betas</span><span class="p">[</span><span class="n">idx_to_change</span><span class="p">,</span> <span class="kc">None</span><span class="p">])</span> <span class="o">*</span> <span class="n">X</span><span class="p">[</span><span class="n">index_permute</span><span class="p">][</span><span class="n">idx_to_change</span><span class="p">]</span><span class="o">.</span><span class="n">view</span><span class="p">(</span><span class="n">X</span><span class="p">[</span><span class="n">idx_to_change</span><span class="p">]</span><span class="o">.</span><span class="n">size</span><span class="p">())</span>  <span class="c1"># noqa</span>
+
+        <span class="k">return</span> <span class="n">X</span><span class="p">,</span> <span class="n">y</span></div>
+</pre></div>
+
+           </div>
+           
+          </div>
+          <footer>
+  
+
+  <hr/>
+
+  <div role="contentinfo">
+    <p>
+        
+        &copy; Copyright 2019, Dreamquark
+
+    </p>
+  </div>
+    
+    
+    
+    Built with <a href="http://sphinx-doc.org/">Sphinx</a> using a
+    
+    <a href="https://github.com/rtfd/sphinx_rtd_theme">theme</a>
+    
+    provided by <a href="https://readthedocs.org">Read the Docs</a>. 
+
+</footer>
+
+        </div>
+      </div>
+
+    </section>
+
+  </div>
+  
+
+  <script type="text/javascript">
+      jQuery(function () {
+          SphinxRtdTheme.Navigation.enable(true);
+      });
+  </script>
+
+  
+  
+    
+   
+
+</body>
+</html>
\ No newline at end of file
diff --git a/docs/_modules/pytorch_tabnet/callbacks.html b/docs/_modules/pytorch_tabnet/callbacks.html
index a59b2516..b4168aaf 100644
--- a/docs/_modules/pytorch_tabnet/callbacks.html
+++ b/docs/_modules/pytorch_tabnet/callbacks.html
@@ -82,14 +82,17 @@
               
             
             
-              <p class="caption"><span class="caption-text">Contents:</span></p>
+              <p><span class="caption-text">Contents:</span></p>
 <ul>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html">README</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#tabnet-attentive-interpretable-tabular-learning">TabNet : Attentive Interpretable Tabular Learning</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#installation">Installation</a></li>
-<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#what-problems-does-pytorch-tabnet-handles">What problems does pytorch-tabnet handles?</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#contributing">Contributing</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#what-problems-does-pytorch-tabnet-handle">What problems does pytorch-tabnet handle?</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#how-to-use-it">How to use it?</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#semi-supervised-pre-training">Semi-supervised pre-training</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#data-augmentation-on-the-fly">Data augmentation on the fly</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#easy-saving-and-loading">Easy saving and loading</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#useful-links">Useful links</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/pytorch_tabnet.html">pytorch_tabnet package</a></li>
 </ul>
@@ -162,6 +165,7 @@ <h1>Source code for pytorch_tabnet.callbacks</h1><div class="highlight"><pre>
 <span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
 <span class="kn">from</span> <span class="nn">dataclasses</span> <span class="kn">import</span> <span class="n">dataclass</span><span class="p">,</span> <span class="n">field</span>
 <span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">List</span><span class="p">,</span> <span class="n">Any</span>
+<span class="kn">import</span> <span class="nn">warnings</span>
 
 
 <div class="viewcode-block" id="Callback"><a class="viewcode-back" href="../../generated_docs/pytorch_tabnet.html#pytorch_tabnet.callbacks.Callback">[docs]</a><span class="k">class</span> <span class="nc">Callback</span><span class="p">:</span>
@@ -325,7 +329,8 @@ <h1>Source code for pytorch_tabnet.callbacks</h1><div class="highlight"><pre>
                 <span class="o">+</span> <span class="sa">f</span><span class="s2">&quot;best_</span><span class="si">{</span><span class="bp">self</span><span class="o">.</span><span class="n">early_stopping_metric</span><span class="si">}</span><span class="s2"> = </span><span class="si">{</span><span class="nb">round</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">best_loss</span><span class="p">,</span> <span class="mi">5</span><span class="p">)</span><span class="si">}</span><span class="s2">&quot;</span>
             <span class="p">)</span>
             <span class="nb">print</span><span class="p">(</span><span class="n">msg</span><span class="p">)</span>
-        <span class="nb">print</span><span class="p">(</span><span class="s2">&quot;Best weights from best epoch are automatically used!&quot;</span><span class="p">)</span></div></div>
+        <span class="n">wrn_msg</span> <span class="o">=</span> <span class="s2">&quot;Best weights from best epoch are automatically used!&quot;</span>
+        <span class="n">warnings</span><span class="o">.</span><span class="n">warn</span><span class="p">(</span><span class="n">wrn_msg</span><span class="p">)</span></div></div>
 
 
 <div class="viewcode-block" id="History"><a class="viewcode-back" href="../../generated_docs/pytorch_tabnet.html#pytorch_tabnet.callbacks.History">[docs]</a><span class="nd">@dataclass</span>
diff --git a/docs/_modules/pytorch_tabnet/metrics.html b/docs/_modules/pytorch_tabnet/metrics.html
index d807616e..267c5532 100644
--- a/docs/_modules/pytorch_tabnet/metrics.html
+++ b/docs/_modules/pytorch_tabnet/metrics.html
@@ -82,14 +82,17 @@
               
             
             
-              <p class="caption"><span class="caption-text">Contents:</span></p>
+              <p><span class="caption-text">Contents:</span></p>
 <ul>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html">README</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#tabnet-attentive-interpretable-tabular-learning">TabNet : Attentive Interpretable Tabular Learning</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#installation">Installation</a></li>
-<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#what-problems-does-pytorch-tabnet-handles">What problems does pytorch-tabnet handles?</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#contributing">Contributing</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#what-problems-does-pytorch-tabnet-handle">What problems does pytorch-tabnet handle?</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#how-to-use-it">How to use it?</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#semi-supervised-pre-training">Semi-supervised pre-training</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#data-augmentation-on-the-fly">Data augmentation on the fly</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#easy-saving-and-loading">Easy saving and loading</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#useful-links">Useful links</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/pytorch_tabnet.html">pytorch_tabnet package</a></li>
 </ul>
@@ -197,7 +200,11 @@ <h1>Source code for pytorch_tabnet.metrics</h1><div class="highlight"><pre>
 <span class="sd">    &quot;&quot;&quot;</span>
     <span class="n">errors</span> <span class="o">=</span> <span class="n">y_pred</span> <span class="o">-</span> <span class="n">embedded_x</span>
     <span class="n">reconstruction_errors</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">mul</span><span class="p">(</span><span class="n">errors</span><span class="p">,</span> <span class="n">obf_vars</span><span class="p">)</span> <span class="o">**</span> <span class="mi">2</span>
-    <span class="n">batch_stds</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">std</span><span class="p">(</span><span class="n">embedded_x</span><span class="p">,</span> <span class="n">dim</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span> <span class="o">**</span> <span class="mi">2</span> <span class="o">+</span> <span class="n">eps</span>
+    <span class="n">batch_means</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">mean</span><span class="p">(</span><span class="n">embedded_x</span><span class="p">,</span> <span class="n">dim</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span>
+    <span class="n">batch_means</span><span class="p">[</span><span class="n">batch_means</span> <span class="o">==</span> <span class="mi">0</span><span class="p">]</span> <span class="o">=</span> <span class="mi">1</span>
+
+    <span class="n">batch_stds</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">std</span><span class="p">(</span><span class="n">embedded_x</span><span class="p">,</span> <span class="n">dim</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span> <span class="o">**</span> <span class="mi">2</span>
+    <span class="n">batch_stds</span><span class="p">[</span><span class="n">batch_stds</span> <span class="o">==</span> <span class="mi">0</span><span class="p">]</span> <span class="o">=</span> <span class="n">batch_means</span><span class="p">[</span><span class="n">batch_stds</span> <span class="o">==</span> <span class="mi">0</span><span class="p">]</span>
     <span class="n">features_loss</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">matmul</span><span class="p">(</span><span class="n">reconstruction_errors</span><span class="p">,</span> <span class="mi">1</span> <span class="o">/</span> <span class="n">batch_stds</span><span class="p">)</span>
     <span class="c1"># compute the number of obfuscated variables to reconstruct</span>
     <span class="n">nb_reconstructed_variables</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">obf_vars</span><span class="p">,</span> <span class="n">dim</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>
@@ -208,6 +215,24 @@ <h1>Source code for pytorch_tabnet.metrics</h1><div class="highlight"><pre>
     <span class="k">return</span> <span class="n">loss</span></div>
 
 
+<div class="viewcode-block" id="UnsupervisedLossNumpy"><a class="viewcode-back" href="../../generated_docs/pytorch_tabnet.html#pytorch_tabnet.metrics.UnsupervisedLossNumpy">[docs]</a><span class="k">def</span> <span class="nf">UnsupervisedLossNumpy</span><span class="p">(</span><span class="n">y_pred</span><span class="p">,</span> <span class="n">embedded_x</span><span class="p">,</span> <span class="n">obf_vars</span><span class="p">,</span> <span class="n">eps</span><span class="o">=</span><span class="mf">1e-9</span><span class="p">):</span>
+    <span class="n">errors</span> <span class="o">=</span> <span class="n">y_pred</span> <span class="o">-</span> <span class="n">embedded_x</span>
+    <span class="n">reconstruction_errors</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">multiply</span><span class="p">(</span><span class="n">errors</span><span class="p">,</span> <span class="n">obf_vars</span><span class="p">)</span> <span class="o">**</span> <span class="mi">2</span>
+    <span class="n">batch_means</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">mean</span><span class="p">(</span><span class="n">embedded_x</span><span class="p">,</span> <span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span>
+    <span class="n">batch_means</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">where</span><span class="p">(</span><span class="n">batch_means</span> <span class="o">==</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="n">batch_means</span><span class="p">)</span>
+
+    <span class="n">batch_stds</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">std</span><span class="p">(</span><span class="n">embedded_x</span><span class="p">,</span> <span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span> <span class="n">ddof</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span> <span class="o">**</span> <span class="mi">2</span>
+    <span class="n">batch_stds</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">where</span><span class="p">(</span><span class="n">batch_stds</span> <span class="o">==</span> <span class="mi">0</span><span class="p">,</span> <span class="n">batch_means</span><span class="p">,</span> <span class="n">batch_stds</span><span class="p">)</span>
+    <span class="n">features_loss</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">matmul</span><span class="p">(</span><span class="n">reconstruction_errors</span><span class="p">,</span> <span class="mi">1</span> <span class="o">/</span> <span class="n">batch_stds</span><span class="p">)</span>
+    <span class="c1"># compute the number of obfuscated variables to reconstruct</span>
+    <span class="n">nb_reconstructed_variables</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">sum</span><span class="p">(</span><span class="n">obf_vars</span><span class="p">,</span> <span class="n">axis</span><span class="o">=</span><span class="mi">1</span><span class="p">)</span>
+    <span class="c1"># take the mean of the reconstructed variable errors</span>
+    <span class="n">features_loss</span> <span class="o">=</span> <span class="n">features_loss</span> <span class="o">/</span> <span class="p">(</span><span class="n">nb_reconstructed_variables</span> <span class="o">+</span> <span class="n">eps</span><span class="p">)</span>
+    <span class="c1"># here we take the mean per batch, contrary to the paper</span>
+    <span class="n">loss</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">mean</span><span class="p">(</span><span class="n">features_loss</span><span class="p">)</span>
+    <span class="k">return</span> <span class="n">loss</span></div>
+
+
 <div class="viewcode-block" id="UnsupMetricContainer"><a class="viewcode-back" href="../../generated_docs/pytorch_tabnet.html#pytorch_tabnet.metrics.UnsupMetricContainer">[docs]</a><span class="nd">@dataclass</span>
 <span class="k">class</span> <span class="nc">UnsupMetricContainer</span><span class="p">:</span>
     <span class="sd">&quot;&quot;&quot;Container holding a list of metrics.</span>
@@ -571,6 +596,41 @@ <h1>Source code for pytorch_tabnet.metrics</h1><div class="highlight"><pre>
         <span class="k">return</span> <span class="n">loss</span><span class="o">.</span><span class="n">item</span><span class="p">()</span></div>
 
 
+<div class="viewcode-block" id="UnsupervisedNumpyMetric"><a class="viewcode-back" href="../../generated_docs/pytorch_tabnet.html#pytorch_tabnet.metrics.UnsupervisedNumpyMetric">[docs]</a><span class="k">class</span> <span class="nc">UnsupervisedNumpyMetric</span><span class="p">(</span><span class="n">Metric</span><span class="p">):</span>
+    <span class="sd">&quot;&quot;&quot;</span>
+<span class="sd">    Unsupervised metric</span>
+<span class="sd">    &quot;&quot;&quot;</span>
+
+    <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span><span class="bp">self</span><span class="p">):</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">_name</span> <span class="o">=</span> <span class="s2">&quot;unsup_loss_numpy&quot;</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">_maximize</span> <span class="o">=</span> <span class="kc">False</span>
+
+    <span class="k">def</span> <span class="fm">__call__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">y_pred</span><span class="p">,</span> <span class="n">embedded_x</span><span class="p">,</span> <span class="n">obf_vars</span><span class="p">):</span>
+        <span class="sd">&quot;&quot;&quot;</span>
+<span class="sd">        Compute MSE (Mean Squared Error) of predictions.</span>
+
+<span class="sd">        Parameters</span>
+<span class="sd">        ----------</span>
+<span class="sd">        y_pred : torch.Tensor or np.array</span>
+<span class="sd">            Reconstructed prediction (with embeddings)</span>
+<span class="sd">        embedded_x : torch.Tensor</span>
+<span class="sd">            Original input embedded by network</span>
+<span class="sd">        obf_vars : torch.Tensor</span>
+<span class="sd">            Binary mask for obfuscated variables.</span>
+<span class="sd">            1 means the variables was obfuscated so reconstruction is based on this.</span>
+
+<span class="sd">        Returns</span>
+<span class="sd">        -------</span>
+<span class="sd">        float</span>
+<span class="sd">            MSE of predictions vs targets.</span>
+<span class="sd">        &quot;&quot;&quot;</span>
+        <span class="k">return</span> <span class="n">UnsupervisedLossNumpy</span><span class="p">(</span>
+            <span class="n">y_pred</span><span class="p">,</span>
+            <span class="n">embedded_x</span><span class="p">,</span>
+            <span class="n">obf_vars</span>
+        <span class="p">)</span></div>
+
+
 <div class="viewcode-block" id="RMSE"><a class="viewcode-back" href="../../generated_docs/pytorch_tabnet.html#pytorch_tabnet.metrics.RMSE">[docs]</a><span class="k">class</span> <span class="nc">RMSE</span><span class="p">(</span><span class="n">Metric</span><span class="p">):</span>
     <span class="sd">&quot;&quot;&quot;</span>
 <span class="sd">    Root Mean Squared Error.</span>
diff --git a/docs/_modules/pytorch_tabnet/multiclass_utils.html b/docs/_modules/pytorch_tabnet/multiclass_utils.html
index e579c068..0a68885b 100644
--- a/docs/_modules/pytorch_tabnet/multiclass_utils.html
+++ b/docs/_modules/pytorch_tabnet/multiclass_utils.html
@@ -82,14 +82,17 @@
               
             
             
-              <p class="caption"><span class="caption-text">Contents:</span></p>
+              <p><span class="caption-text">Contents:</span></p>
 <ul>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html">README</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#tabnet-attentive-interpretable-tabular-learning">TabNet : Attentive Interpretable Tabular Learning</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#installation">Installation</a></li>
-<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#what-problems-does-pytorch-tabnet-handles">What problems does pytorch-tabnet handles?</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#contributing">Contributing</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#what-problems-does-pytorch-tabnet-handle">What problems does pytorch-tabnet handle?</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#how-to-use-it">How to use it?</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#semi-supervised-pre-training">Semi-supervised pre-training</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#data-augmentation-on-the-fly">Data augmentation on the fly</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#easy-saving-and-loading">Easy saving and loading</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#useful-links">Useful links</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/pytorch_tabnet.html">pytorch_tabnet package</a></li>
 </ul>
diff --git a/docs/_modules/pytorch_tabnet/multitask.html b/docs/_modules/pytorch_tabnet/multitask.html
index 067feb01..e16ee2e6 100644
--- a/docs/_modules/pytorch_tabnet/multitask.html
+++ b/docs/_modules/pytorch_tabnet/multitask.html
@@ -82,14 +82,17 @@
               
             
             
-              <p class="caption"><span class="caption-text">Contents:</span></p>
+              <p><span class="caption-text">Contents:</span></p>
 <ul>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html">README</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#tabnet-attentive-interpretable-tabular-learning">TabNet : Attentive Interpretable Tabular Learning</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#installation">Installation</a></li>
-<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#what-problems-does-pytorch-tabnet-handles">What problems does pytorch-tabnet handles?</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#contributing">Contributing</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#what-problems-does-pytorch-tabnet-handle">What problems does pytorch-tabnet handle?</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#how-to-use-it">How to use it?</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#semi-supervised-pre-training">Semi-supervised pre-training</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#data-augmentation-on-the-fly">Data augmentation on the fly</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#easy-saving-and-loading">Easy saving and loading</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#useful-links">Useful links</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/pytorch_tabnet.html">pytorch_tabnet package</a></li>
 </ul>
diff --git a/docs/_modules/pytorch_tabnet/pretraining.html b/docs/_modules/pytorch_tabnet/pretraining.html
index 9244752b..a4f95e74 100644
--- a/docs/_modules/pytorch_tabnet/pretraining.html
+++ b/docs/_modules/pytorch_tabnet/pretraining.html
@@ -82,14 +82,17 @@
               
             
             
-              <p class="caption"><span class="caption-text">Contents:</span></p>
+              <p><span class="caption-text">Contents:</span></p>
 <ul>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html">README</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#tabnet-attentive-interpretable-tabular-learning">TabNet : Attentive Interpretable Tabular Learning</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#installation">Installation</a></li>
-<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#what-problems-does-pytorch-tabnet-handles">What problems does pytorch-tabnet handles?</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#contributing">Contributing</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#what-problems-does-pytorch-tabnet-handle">What problems does pytorch-tabnet handle?</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#how-to-use-it">How to use it?</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#semi-supervised-pre-training">Semi-supervised pre-training</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#data-augmentation-on-the-fly">Data augmentation on the fly</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#easy-saving-and-loading">Easy saving and loading</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#useful-links">Useful links</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/pytorch_tabnet.html">pytorch_tabnet package</a></li>
 </ul>
@@ -158,13 +161,13 @@
   <h1>Source code for pytorch_tabnet.pretraining</h1><div class="highlight"><pre>
 <span></span><span class="kn">import</span> <span class="nn">torch</span>
 <span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
-<span class="kn">from</span> <span class="nn">sklearn.utils</span> <span class="kn">import</span> <span class="n">check_array</span>
 <span class="kn">from</span> <span class="nn">torch.utils.data</span> <span class="kn">import</span> <span class="n">DataLoader</span>
 <span class="kn">from</span> <span class="nn">pytorch_tabnet</span> <span class="kn">import</span> <span class="n">tab_network</span>
 <span class="kn">from</span> <span class="nn">pytorch_tabnet.utils</span> <span class="kn">import</span> <span class="p">(</span>
     <span class="n">create_explain_matrix</span><span class="p">,</span>
     <span class="n">filter_weights</span><span class="p">,</span>
-    <span class="n">PredictDataset</span>
+    <span class="n">PredictDataset</span><span class="p">,</span>
+    <span class="n">check_input</span>
 <span class="p">)</span>
 <span class="kn">from</span> <span class="nn">torch.nn.utils</span> <span class="kn">import</span> <span class="n">clip_grad_norm_</span>
 <span class="kn">from</span> <span class="nn">pytorch_tabnet.pretraining_utils</span> <span class="kn">import</span> <span class="p">(</span>
@@ -184,7 +187,7 @@ <h1>Source code for pytorch_tabnet.pretraining</h1><div class="highlight"><pre>
         <span class="nb">super</span><span class="p">(</span><span class="n">TabNetPretrainer</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="n">__post_init__</span><span class="p">()</span>
         <span class="bp">self</span><span class="o">.</span><span class="n">_task</span> <span class="o">=</span> <span class="s1">&#39;unsupervised&#39;</span>
         <span class="bp">self</span><span class="o">.</span><span class="n">_default_loss</span> <span class="o">=</span> <span class="n">UnsupervisedLoss</span>
-        <span class="bp">self</span><span class="o">.</span><span class="n">_default_metric</span> <span class="o">=</span> <span class="s1">&#39;unsup_loss&#39;</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">_default_metric</span> <span class="o">=</span> <span class="s1">&#39;unsup_loss_numpy&#39;</span>
 
 <div class="viewcode-block" id="TabNetPretrainer.prepare_target"><a class="viewcode-back" href="../../generated_docs/pytorch_tabnet.html#pytorch_tabnet.pretraining.TabNetPretrainer.prepare_target">[docs]</a>    <span class="k">def</span> <span class="nf">prepare_target</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">y</span><span class="p">):</span>
         <span class="k">return</span> <span class="n">y</span></div>
@@ -213,9 +216,10 @@ <h1>Source code for pytorch_tabnet.pretraining</h1><div class="highlight"><pre>
         <span class="n">batch_size</span><span class="o">=</span><span class="mi">1024</span><span class="p">,</span>
         <span class="n">virtual_batch_size</span><span class="o">=</span><span class="mi">128</span><span class="p">,</span>
         <span class="n">num_workers</span><span class="o">=</span><span class="mi">0</span><span class="p">,</span>
-        <span class="n">drop_last</span><span class="o">=</span><span class="kc">False</span><span class="p">,</span>
+        <span class="n">drop_last</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
         <span class="n">callbacks</span><span class="o">=</span><span class="kc">None</span><span class="p">,</span>
         <span class="n">pin_memory</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span>
+        <span class="n">warm_start</span><span class="o">=</span><span class="kc">False</span>
     <span class="p">):</span>
         <span class="sd">&quot;&quot;&quot;Train a neural network stored in self.network</span>
 <span class="sd">        Using train_dataloader for training data and</span>
@@ -276,7 +280,7 @@ <h1>Source code for pytorch_tabnet.pretraining</h1><div class="highlight"><pre>
         <span class="k">else</span><span class="p">:</span>
             <span class="bp">self</span><span class="o">.</span><span class="n">loss_fn</span> <span class="o">=</span> <span class="n">loss_fn</span>
 
-        <span class="n">check_array</span><span class="p">(</span><span class="n">X_train</span><span class="p">)</span>
+        <span class="n">check_input</span><span class="p">(</span><span class="n">X_train</span><span class="p">)</span>
 
         <span class="bp">self</span><span class="o">.</span><span class="n">update_fit_params</span><span class="p">(</span>
             <span class="n">weights</span><span class="p">,</span>
@@ -288,8 +292,10 @@ <h1>Source code for pytorch_tabnet.pretraining</h1><div class="highlight"><pre>
             <span class="n">X_train</span><span class="p">,</span> <span class="n">eval_set</span>
         <span class="p">)</span>
 
-        <span class="k">if</span> <span class="ow">not</span> <span class="nb">hasattr</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="s1">&#39;network&#39;</span><span class="p">):</span>
+        <span class="k">if</span> <span class="ow">not</span> <span class="nb">hasattr</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="s2">&quot;network&quot;</span><span class="p">)</span> <span class="ow">or</span> <span class="ow">not</span> <span class="n">warm_start</span><span class="p">:</span>
+            <span class="c1"># model has never been fitted before of warm_start is False</span>
             <span class="bp">self</span><span class="o">.</span><span class="n">_set_network</span><span class="p">()</span>
+
         <span class="bp">self</span><span class="o">.</span><span class="n">_update_network_params</span><span class="p">()</span>
         <span class="bp">self</span><span class="o">.</span><span class="n">_set_metrics</span><span class="p">(</span><span class="n">eval_names</span><span class="p">)</span>
         <span class="bp">self</span><span class="o">.</span><span class="n">_set_optimizer</span><span class="p">()</span>
@@ -326,6 +332,7 @@ <h1>Source code for pytorch_tabnet.pretraining</h1><div class="highlight"><pre>
         <span class="sd">&quot;&quot;&quot;Setup the network and explain matrix.&quot;&quot;&quot;</span>
         <span class="k">if</span> <span class="ow">not</span> <span class="nb">hasattr</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="s1">&#39;pretraining_ratio&#39;</span><span class="p">):</span>
             <span class="bp">self</span><span class="o">.</span><span class="n">pretraining_ratio</span> <span class="o">=</span> <span class="mf">0.5</span>
+        <span class="n">torch</span><span class="o">.</span><span class="n">manual_seed</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">seed</span><span class="p">)</span>
         <span class="bp">self</span><span class="o">.</span><span class="n">network</span> <span class="o">=</span> <span class="n">tab_network</span><span class="o">.</span><span class="n">TabNetPretraining</span><span class="p">(</span>
             <span class="bp">self</span><span class="o">.</span><span class="n">input_dim</span><span class="p">,</span>
             <span class="n">pretraining_ratio</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">pretraining_ratio</span><span class="p">,</span>
@@ -495,9 +502,9 @@ <h1>Source code for pytorch_tabnet.pretraining</h1><div class="highlight"><pre>
         <span class="c1"># Main loop</span>
         <span class="k">for</span> <span class="n">batch_idx</span><span class="p">,</span> <span class="n">X</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">loader</span><span class="p">):</span>
             <span class="n">output</span><span class="p">,</span> <span class="n">embedded_x</span><span class="p">,</span> <span class="n">obf_vars</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">_predict_batch</span><span class="p">(</span><span class="n">X</span><span class="p">)</span>
-            <span class="n">list_output</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">output</span><span class="p">)</span>
-            <span class="n">list_embedded_x</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">embedded_x</span><span class="p">)</span>
-            <span class="n">list_obfuscation</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">obf_vars</span><span class="p">)</span>
+            <span class="n">list_output</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">output</span><span class="o">.</span><span class="n">cpu</span><span class="p">()</span><span class="o">.</span><span class="n">detach</span><span class="p">()</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span>
+            <span class="n">list_embedded_x</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">embedded_x</span><span class="o">.</span><span class="n">cpu</span><span class="p">()</span><span class="o">.</span><span class="n">detach</span><span class="p">()</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span>
+            <span class="n">list_obfuscation</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">obf_vars</span><span class="o">.</span><span class="n">cpu</span><span class="p">()</span><span class="o">.</span><span class="n">detach</span><span class="p">()</span><span class="o">.</span><span class="n">numpy</span><span class="p">())</span>
 
         <span class="n">output</span><span class="p">,</span> <span class="n">embedded_x</span><span class="p">,</span> <span class="n">obf_vars</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">stack_batches</span><span class="p">(</span><span class="n">list_output</span><span class="p">,</span>
                                                           <span class="n">list_embedded_x</span><span class="p">,</span>
@@ -526,9 +533,9 @@ <h1>Source code for pytorch_tabnet.pretraining</h1><div class="highlight"><pre>
         <span class="k">return</span> <span class="bp">self</span><span class="o">.</span><span class="n">network</span><span class="p">(</span><span class="n">X</span><span class="p">)</span>
 
 <div class="viewcode-block" id="TabNetPretrainer.stack_batches"><a class="viewcode-back" href="../../generated_docs/pytorch_tabnet.html#pytorch_tabnet.pretraining.TabNetPretrainer.stack_batches">[docs]</a>    <span class="k">def</span> <span class="nf">stack_batches</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">list_output</span><span class="p">,</span> <span class="n">list_embedded_x</span><span class="p">,</span> <span class="n">list_obfuscation</span><span class="p">):</span>
-        <span class="n">output</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">(</span><span class="n">list_output</span><span class="p">,</span> <span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span>
-        <span class="n">embedded_x</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">(</span><span class="n">list_embedded_x</span><span class="p">,</span> <span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span>
-        <span class="n">obf_vars</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">cat</span><span class="p">(</span><span class="n">list_obfuscation</span><span class="p">,</span> <span class="n">axis</span><span class="o">=</span><span class="mi">0</span><span class="p">)</span>
+        <span class="n">output</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">vstack</span><span class="p">(</span><span class="n">list_output</span><span class="p">)</span>
+        <span class="n">embedded_x</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">vstack</span><span class="p">(</span><span class="n">list_embedded_x</span><span class="p">)</span>
+        <span class="n">obf_vars</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">vstack</span><span class="p">(</span><span class="n">list_obfuscation</span><span class="p">)</span>
         <span class="k">return</span> <span class="n">output</span><span class="p">,</span> <span class="n">embedded_x</span><span class="p">,</span> <span class="n">obf_vars</span></div>
 
 <div class="viewcode-block" id="TabNetPretrainer.predict"><a class="viewcode-back" href="../../generated_docs/pytorch_tabnet.html#pytorch_tabnet.pretraining.TabNetPretrainer.predict">[docs]</a>    <span class="k">def</span> <span class="nf">predict</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">X</span><span class="p">):</span>
diff --git a/docs/_modules/pytorch_tabnet/pretraining_utils.html b/docs/_modules/pytorch_tabnet/pretraining_utils.html
index 43c1af09..1234f03b 100644
--- a/docs/_modules/pytorch_tabnet/pretraining_utils.html
+++ b/docs/_modules/pytorch_tabnet/pretraining_utils.html
@@ -82,14 +82,17 @@
               
             
             
-              <p class="caption"><span class="caption-text">Contents:</span></p>
+              <p><span class="caption-text">Contents:</span></p>
 <ul>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html">README</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#tabnet-attentive-interpretable-tabular-learning">TabNet : Attentive Interpretable Tabular Learning</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#installation">Installation</a></li>
-<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#what-problems-does-pytorch-tabnet-handles">What problems does pytorch-tabnet handles?</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#contributing">Contributing</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#what-problems-does-pytorch-tabnet-handle">What problems does pytorch-tabnet handle?</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#how-to-use-it">How to use it?</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#semi-supervised-pre-training">Semi-supervised pre-training</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#data-augmentation-on-the-fly">Data augmentation on the fly</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#easy-saving-and-loading">Easy saving and loading</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#useful-links">Useful links</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/pytorch_tabnet.html">pytorch_tabnet package</a></li>
 </ul>
@@ -160,8 +163,8 @@ <h1>Source code for pytorch_tabnet.pretraining_utils</h1><div class="highlight">
 <span class="kn">from</span> <span class="nn">pytorch_tabnet.utils</span> <span class="kn">import</span> <span class="p">(</span>
     <span class="n">create_sampler</span><span class="p">,</span>
     <span class="n">PredictDataset</span><span class="p">,</span>
+    <span class="n">check_input</span>
 <span class="p">)</span>
-<span class="kn">from</span> <span class="nn">sklearn.utils</span> <span class="kn">import</span> <span class="n">check_array</span>
 
 
 <div class="viewcode-block" id="create_dataloaders"><a class="viewcode-back" href="../../generated_docs/pytorch_tabnet.html#pytorch_tabnet.pretraining_utils.create_dataloaders">[docs]</a><span class="k">def</span> <span class="nf">create_dataloaders</span><span class="p">(</span>
@@ -251,7 +254,7 @@ <h1>Source code for pytorch_tabnet.pretraining_utils</h1><div class="highlight">
     <span class="p">),</span> <span class="s2">&quot;eval_set and eval_name have not the same length&quot;</span>
 
     <span class="k">for</span> <span class="n">set_nb</span><span class="p">,</span> <span class="n">X</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">eval_set</span><span class="p">):</span>
-        <span class="n">check_array</span><span class="p">(</span><span class="n">X</span><span class="p">)</span>
+        <span class="n">check_input</span><span class="p">(</span><span class="n">X</span><span class="p">)</span>
         <span class="n">msg</span> <span class="o">=</span> <span class="p">(</span>
             <span class="sa">f</span><span class="s2">&quot;Number of columns is different between eval set </span><span class="si">{</span><span class="n">set_nb</span><span class="si">}</span><span class="s2">&quot;</span>
             <span class="o">+</span> <span class="sa">f</span><span class="s2">&quot;(</span><span class="si">{</span><span class="n">X</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span><span class="si">}</span><span class="s2">) and X_train (</span><span class="si">{</span><span class="n">X_train</span><span class="o">.</span><span class="n">shape</span><span class="p">[</span><span class="mi">1</span><span class="p">]</span><span class="si">}</span><span class="s2">)&quot;</span>
diff --git a/docs/_modules/pytorch_tabnet/sparsemax.html b/docs/_modules/pytorch_tabnet/sparsemax.html
index 0595e900..8f03e2ea 100644
--- a/docs/_modules/pytorch_tabnet/sparsemax.html
+++ b/docs/_modules/pytorch_tabnet/sparsemax.html
@@ -82,14 +82,17 @@
               
             
             
-              <p class="caption"><span class="caption-text">Contents:</span></p>
+              <p><span class="caption-text">Contents:</span></p>
 <ul>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html">README</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#tabnet-attentive-interpretable-tabular-learning">TabNet : Attentive Interpretable Tabular Learning</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#installation">Installation</a></li>
-<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#what-problems-does-pytorch-tabnet-handles">What problems does pytorch-tabnet handles?</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#contributing">Contributing</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#what-problems-does-pytorch-tabnet-handle">What problems does pytorch-tabnet handle?</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#how-to-use-it">How to use it?</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#semi-supervised-pre-training">Semi-supervised pre-training</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#data-augmentation-on-the-fly">Data augmentation on the fly</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#easy-saving-and-loading">Easy saving and loading</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#useful-links">Useful links</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/pytorch_tabnet.html">pytorch_tabnet package</a></li>
 </ul>
diff --git a/docs/_modules/pytorch_tabnet/tab_model.html b/docs/_modules/pytorch_tabnet/tab_model.html
index 4c926370..d8bf5379 100644
--- a/docs/_modules/pytorch_tabnet/tab_model.html
+++ b/docs/_modules/pytorch_tabnet/tab_model.html
@@ -82,14 +82,17 @@
               
             
             
-              <p class="caption"><span class="caption-text">Contents:</span></p>
+              <p><span class="caption-text">Contents:</span></p>
 <ul>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html">README</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#tabnet-attentive-interpretable-tabular-learning">TabNet : Attentive Interpretable Tabular Learning</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#installation">Installation</a></li>
-<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#what-problems-does-pytorch-tabnet-handles">What problems does pytorch-tabnet handles?</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#contributing">Contributing</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#what-problems-does-pytorch-tabnet-handle">What problems does pytorch-tabnet handle?</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#how-to-use-it">How to use it?</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#semi-supervised-pre-training">Semi-supervised pre-training</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#data-augmentation-on-the-fly">Data augmentation on the fly</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#easy-saving-and-loading">Easy saving and loading</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#useful-links">Useful links</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/pytorch_tabnet.html">pytorch_tabnet package</a></li>
 </ul>
diff --git a/docs/_modules/pytorch_tabnet/tab_network.html b/docs/_modules/pytorch_tabnet/tab_network.html
index 81f9d77e..c8ba7358 100644
--- a/docs/_modules/pytorch_tabnet/tab_network.html
+++ b/docs/_modules/pytorch_tabnet/tab_network.html
@@ -82,14 +82,17 @@
               
             
             
-              <p class="caption"><span class="caption-text">Contents:</span></p>
+              <p><span class="caption-text">Contents:</span></p>
 <ul>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html">README</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#tabnet-attentive-interpretable-tabular-learning">TabNet : Attentive Interpretable Tabular Learning</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#installation">Installation</a></li>
-<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#what-problems-does-pytorch-tabnet-handles">What problems does pytorch-tabnet handles?</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#contributing">Contributing</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#what-problems-does-pytorch-tabnet-handle">What problems does pytorch-tabnet handle?</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#how-to-use-it">How to use it?</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#semi-supervised-pre-training">Semi-supervised pre-training</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#data-augmentation-on-the-fly">Data augmentation on the fly</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#easy-saving-and-loading">Easy saving and loading</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#useful-links">Useful links</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/pytorch_tabnet.html">pytorch_tabnet package</a></li>
 </ul>
@@ -364,8 +367,8 @@ <h1>Source code for pytorch_tabnet.tab_network</h1><div class="highlight"><pre>
         <span class="n">input_dim</span><span class="p">,</span>
         <span class="n">n_d</span><span class="o">=</span><span class="mi">8</span><span class="p">,</span>
         <span class="n">n_steps</span><span class="o">=</span><span class="mi">3</span><span class="p">,</span>
-        <span class="n">n_independent</span><span class="o">=</span><span class="mi">2</span><span class="p">,</span>
-        <span class="n">n_shared</span><span class="o">=</span><span class="mi">2</span><span class="p">,</span>
+        <span class="n">n_independent</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
+        <span class="n">n_shared</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
         <span class="n">virtual_batch_size</span><span class="o">=</span><span class="mi">128</span><span class="p">,</span>
         <span class="n">momentum</span><span class="o">=</span><span class="mf">0.02</span><span class="p">,</span>
     <span class="p">):</span>
@@ -386,9 +389,9 @@ <h1>Source code for pytorch_tabnet.tab_network</h1><div class="highlight"><pre>
 <span class="sd">        gamma : float</span>
 <span class="sd">            Float above 1, scaling factor for attention updates (usually between 1.0 to 2.0)</span>
 <span class="sd">        n_independent : int</span>
-<span class="sd">            Number of independent GLU layer in each GLU block (default 2)</span>
+<span class="sd">            Number of independent GLU layer in each GLU block (default 1)</span>
 <span class="sd">        n_shared : int</span>
-<span class="sd">            Number of independent GLU layer in each GLU block (default 2)</span>
+<span class="sd">            Number of independent GLU layer in each GLU block (default 1)</span>
 <span class="sd">        virtual_batch_size : int</span>
 <span class="sd">            Batch size for Ghost Batch Normalization</span>
 <span class="sd">        momentum : float</span>
@@ -403,7 +406,6 @@ <h1>Source code for pytorch_tabnet.tab_network</h1><div class="highlight"><pre>
         <span class="bp">self</span><span class="o">.</span><span class="n">virtual_batch_size</span> <span class="o">=</span> <span class="n">virtual_batch_size</span>
 
         <span class="bp">self</span><span class="o">.</span><span class="n">feat_transformers</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">ModuleList</span><span class="p">()</span>
-        <span class="bp">self</span><span class="o">.</span><span class="n">reconstruction_layers</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">ModuleList</span><span class="p">()</span>
 
         <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">n_shared</span> <span class="o">&gt;</span> <span class="mi">0</span><span class="p">:</span>
             <span class="n">shared_feat_transform</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">nn</span><span class="o">.</span><span class="n">ModuleList</span><span class="p">()</span>
@@ -426,16 +428,16 @@ <h1>Source code for pytorch_tabnet.tab_network</h1><div class="highlight"><pre>
                 <span class="n">momentum</span><span class="o">=</span><span class="n">momentum</span><span class="p">,</span>
             <span class="p">)</span>
             <span class="bp">self</span><span class="o">.</span><span class="n">feat_transformers</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">transformer</span><span class="p">)</span>
-            <span class="n">reconstruction_layer</span> <span class="o">=</span> <span class="n">Linear</span><span class="p">(</span><span class="n">n_d</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">input_dim</span><span class="p">,</span> <span class="n">bias</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
-            <span class="n">initialize_non_glu</span><span class="p">(</span><span class="n">reconstruction_layer</span><span class="p">,</span> <span class="n">n_d</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">input_dim</span><span class="p">)</span>
-            <span class="bp">self</span><span class="o">.</span><span class="n">reconstruction_layers</span><span class="o">.</span><span class="n">append</span><span class="p">(</span><span class="n">reconstruction_layer</span><span class="p">)</span>
+
+        <span class="bp">self</span><span class="o">.</span><span class="n">reconstruction_layer</span> <span class="o">=</span> <span class="n">Linear</span><span class="p">(</span><span class="n">n_d</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">input_dim</span><span class="p">,</span> <span class="n">bias</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span>
+        <span class="n">initialize_non_glu</span><span class="p">(</span><span class="bp">self</span><span class="o">.</span><span class="n">reconstruction_layer</span><span class="p">,</span> <span class="n">n_d</span><span class="p">,</span> <span class="bp">self</span><span class="o">.</span><span class="n">input_dim</span><span class="p">)</span>
 
 <div class="viewcode-block" id="TabNetDecoder.forward"><a class="viewcode-back" href="../../generated_docs/pytorch_tabnet.html#pytorch_tabnet.tab_network.TabNetDecoder.forward">[docs]</a>    <span class="k">def</span> <span class="nf">forward</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">steps_output</span><span class="p">):</span>
         <span class="n">res</span> <span class="o">=</span> <span class="mi">0</span>
         <span class="k">for</span> <span class="n">step_nb</span><span class="p">,</span> <span class="n">step_output</span> <span class="ow">in</span> <span class="nb">enumerate</span><span class="p">(</span><span class="n">steps_output</span><span class="p">):</span>
             <span class="n">x</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">feat_transformers</span><span class="p">[</span><span class="n">step_nb</span><span class="p">](</span><span class="n">step_output</span><span class="p">)</span>
-            <span class="n">x</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">reconstruction_layers</span><span class="p">[</span><span class="n">step_nb</span><span class="p">](</span><span class="n">step_output</span><span class="p">)</span>
             <span class="n">res</span> <span class="o">=</span> <span class="n">torch</span><span class="o">.</span><span class="n">add</span><span class="p">(</span><span class="n">res</span><span class="p">,</span> <span class="n">x</span><span class="p">)</span>
+        <span class="n">res</span> <span class="o">=</span> <span class="bp">self</span><span class="o">.</span><span class="n">reconstruction_layer</span><span class="p">(</span><span class="n">res</span><span class="p">)</span>
         <span class="k">return</span> <span class="n">res</span></div></div>
 
 
@@ -457,6 +459,8 @@ <h1>Source code for pytorch_tabnet.tab_network</h1><div class="highlight"><pre>
         <span class="n">virtual_batch_size</span><span class="o">=</span><span class="mi">128</span><span class="p">,</span>
         <span class="n">momentum</span><span class="o">=</span><span class="mf">0.02</span><span class="p">,</span>
         <span class="n">mask_type</span><span class="o">=</span><span class="s2">&quot;sparsemax&quot;</span><span class="p">,</span>
+        <span class="n">n_shared_decoder</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
+        <span class="n">n_indep_decoder</span><span class="o">=</span><span class="mi">1</span><span class="p">,</span>
     <span class="p">):</span>
         <span class="nb">super</span><span class="p">(</span><span class="n">TabNetPretraining</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
 
@@ -474,6 +478,8 @@ <h1>Source code for pytorch_tabnet.tab_network</h1><div class="highlight"><pre>
         <span class="bp">self</span><span class="o">.</span><span class="n">n_shared</span> <span class="o">=</span> <span class="n">n_shared</span>
         <span class="bp">self</span><span class="o">.</span><span class="n">mask_type</span> <span class="o">=</span> <span class="n">mask_type</span>
         <span class="bp">self</span><span class="o">.</span><span class="n">pretraining_ratio</span> <span class="o">=</span> <span class="n">pretraining_ratio</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">n_shared_decoder</span> <span class="o">=</span> <span class="n">n_shared_decoder</span>
+        <span class="bp">self</span><span class="o">.</span><span class="n">n_indep_decoder</span> <span class="o">=</span> <span class="n">n_indep_decoder</span>
 
         <span class="k">if</span> <span class="bp">self</span><span class="o">.</span><span class="n">n_steps</span> <span class="o">&lt;=</span> <span class="mi">0</span><span class="p">:</span>
             <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="s2">&quot;n_steps should be a positive integer.&quot;</span><span class="p">)</span>
@@ -503,8 +509,8 @@ <h1>Source code for pytorch_tabnet.tab_network</h1><div class="highlight"><pre>
             <span class="bp">self</span><span class="o">.</span><span class="n">post_embed_dim</span><span class="p">,</span>
             <span class="n">n_d</span><span class="o">=</span><span class="n">n_d</span><span class="p">,</span>
             <span class="n">n_steps</span><span class="o">=</span><span class="n">n_steps</span><span class="p">,</span>
-            <span class="n">n_independent</span><span class="o">=</span><span class="n">n_independent</span><span class="p">,</span>
-            <span class="n">n_shared</span><span class="o">=</span><span class="n">n_shared</span><span class="p">,</span>
+            <span class="n">n_independent</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">n_indep_decoder</span><span class="p">,</span>
+            <span class="n">n_shared</span><span class="o">=</span><span class="bp">self</span><span class="o">.</span><span class="n">n_shared_decoder</span><span class="p">,</span>
             <span class="n">virtual_batch_size</span><span class="o">=</span><span class="n">virtual_batch_size</span><span class="p">,</span>
             <span class="n">momentum</span><span class="o">=</span><span class="n">momentum</span><span class="p">,</span>
         <span class="p">)</span>
@@ -952,10 +958,19 @@ <h1>Source code for pytorch_tabnet.tab_network</h1><div class="highlight"><pre>
 <span class="sd">            If int, the same embedding dimension will be used for all categorical features</span>
 <span class="sd">        &quot;&quot;&quot;</span>
         <span class="nb">super</span><span class="p">(</span><span class="n">EmbeddingGenerator</span><span class="p">,</span> <span class="bp">self</span><span class="p">)</span><span class="o">.</span><span class="fm">__init__</span><span class="p">()</span>
-        <span class="k">if</span> <span class="n">cat_dims</span> <span class="o">==</span> <span class="p">[]</span> <span class="ow">or</span> <span class="n">cat_idxs</span> <span class="o">==</span> <span class="p">[]:</span>
+        <span class="k">if</span> <span class="n">cat_dims</span> <span class="o">==</span> <span class="p">[]</span> <span class="ow">and</span> <span class="n">cat_idxs</span> <span class="o">==</span> <span class="p">[]:</span>
             <span class="bp">self</span><span class="o">.</span><span class="n">skip_embedding</span> <span class="o">=</span> <span class="kc">True</span>
             <span class="bp">self</span><span class="o">.</span><span class="n">post_embed_dim</span> <span class="o">=</span> <span class="n">input_dim</span>
             <span class="k">return</span>
+        <span class="k">elif</span> <span class="p">(</span><span class="n">cat_dims</span> <span class="o">==</span> <span class="p">[])</span> <span class="o">^</span> <span class="p">(</span><span class="n">cat_idxs</span> <span class="o">==</span> <span class="p">[]):</span>
+            <span class="k">if</span> <span class="n">cat_dims</span> <span class="o">==</span> <span class="p">[]:</span>
+                <span class="n">msg</span> <span class="o">=</span> <span class="s2">&quot;If cat_idxs is non-empty, cat_dims must be defined as a list of same length.&quot;</span>
+            <span class="k">else</span><span class="p">:</span>
+                <span class="n">msg</span> <span class="o">=</span> <span class="s2">&quot;If cat_dims is non-empty, cat_idxs must be defined as a list of same length.&quot;</span>
+            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="n">msg</span><span class="p">)</span>
+        <span class="k">elif</span> <span class="nb">len</span><span class="p">(</span><span class="n">cat_dims</span><span class="p">)</span> <span class="o">!=</span> <span class="nb">len</span><span class="p">(</span><span class="n">cat_idxs</span><span class="p">):</span>
+            <span class="n">msg</span> <span class="o">=</span> <span class="s2">&quot;The lists cat_dims and cat_idxs must have the same length.&quot;</span>
+            <span class="k">raise</span> <span class="ne">ValueError</span><span class="p">(</span><span class="n">msg</span><span class="p">)</span>
 
         <span class="bp">self</span><span class="o">.</span><span class="n">skip_embedding</span> <span class="o">=</span> <span class="kc">False</span>
         <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">cat_emb_dim</span><span class="p">,</span> <span class="nb">int</span><span class="p">):</span>
diff --git a/docs/_modules/pytorch_tabnet/utils.html b/docs/_modules/pytorch_tabnet/utils.html
index a9c18e4f..17205ca8 100644
--- a/docs/_modules/pytorch_tabnet/utils.html
+++ b/docs/_modules/pytorch_tabnet/utils.html
@@ -82,14 +82,17 @@
               
             
             
-              <p class="caption"><span class="caption-text">Contents:</span></p>
+              <p><span class="caption-text">Contents:</span></p>
 <ul>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html">README</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#tabnet-attentive-interpretable-tabular-learning">TabNet : Attentive Interpretable Tabular Learning</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#installation">Installation</a></li>
-<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#what-problems-does-pytorch-tabnet-handles">What problems does pytorch-tabnet handles?</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#contributing">Contributing</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#what-problems-does-pytorch-tabnet-handle">What problems does pytorch-tabnet handle?</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#how-to-use-it">How to use it?</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#semi-supervised-pre-training">Semi-supervised pre-training</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#data-augmentation-on-the-fly">Data augmentation on the fly</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#easy-saving-and-loading">Easy saving and loading</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/README.html#useful-links">Useful links</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../generated_docs/pytorch_tabnet.html">pytorch_tabnet package</a></li>
 </ul>
@@ -163,6 +166,8 @@ <h1>Source code for pytorch_tabnet.utils</h1><div class="highlight"><pre>
 <span class="kn">import</span> <span class="nn">scipy</span>
 <span class="kn">import</span> <span class="nn">json</span>
 <span class="kn">from</span> <span class="nn">sklearn.utils</span> <span class="kn">import</span> <span class="n">check_array</span>
+<span class="kn">import</span> <span class="nn">pandas</span> <span class="k">as</span> <span class="nn">pd</span>
+<span class="kn">import</span> <span class="nn">warnings</span>
 
 
 <div class="viewcode-block" id="TorchDataset"><a class="viewcode-back" href="../../generated_docs/pytorch_tabnet.html#pytorch_tabnet.utils.TorchDataset">[docs]</a><span class="k">class</span> <span class="nc">TorchDataset</span><span class="p">(</span><span class="n">Dataset</span><span class="p">):</span>
@@ -429,7 +434,7 @@ <h1>Source code for pytorch_tabnet.utils</h1><div class="highlight"><pre>
             <span class="nb">len</span><span class="p">(</span><span class="n">elem</span><span class="p">)</span> <span class="o">==</span> <span class="mi">2</span> <span class="k">for</span> <span class="n">elem</span> <span class="ow">in</span> <span class="n">eval_set</span>
         <span class="p">),</span> <span class="s2">&quot;Each tuple of eval_set need to have two elements&quot;</span>
     <span class="k">for</span> <span class="n">name</span><span class="p">,</span> <span class="p">(</span><span class="n">X</span><span class="p">,</span> <span class="n">y</span><span class="p">)</span> <span class="ow">in</span> <span class="nb">zip</span><span class="p">(</span><span class="n">eval_name</span><span class="p">,</span> <span class="n">eval_set</span><span class="p">):</span>
-        <span class="n">check_array</span><span class="p">(</span><span class="n">X</span><span class="p">)</span>
+        <span class="n">check_input</span><span class="p">(</span><span class="n">X</span><span class="p">)</span>
         <span class="n">msg</span> <span class="o">=</span> <span class="p">(</span>
             <span class="sa">f</span><span class="s2">&quot;Dimension mismatch between X_</span><span class="si">{</span><span class="n">name</span><span class="si">}</span><span class="s2"> &quot;</span>
             <span class="o">+</span> <span class="sa">f</span><span class="s2">&quot;</span><span class="si">{</span><span class="n">X</span><span class="o">.</span><span class="n">shape</span><span class="si">}</span><span class="s2"> and X_train </span><span class="si">{</span><span class="n">X_train</span><span class="o">.</span><span class="n">shape</span><span class="si">}</span><span class="s2">&quot;</span>
@@ -495,6 +500,28 @@ <h1>Source code for pytorch_tabnet.utils</h1><div class="highlight"><pre>
             <span class="k">return</span> <span class="nb">int</span><span class="p">(</span><span class="n">obj</span><span class="p">)</span>
         <span class="c1"># Let the base class default method raise the TypeError</span>
         <span class="k">return</span> <span class="n">json</span><span class="o">.</span><span class="n">JSONEncoder</span><span class="o">.</span><span class="n">default</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">obj</span><span class="p">)</span></div></div>
+
+
+<div class="viewcode-block" id="check_input"><a class="viewcode-back" href="../../generated_docs/pytorch_tabnet.html#pytorch_tabnet.utils.check_input">[docs]</a><span class="k">def</span> <span class="nf">check_input</span><span class="p">(</span><span class="n">X</span><span class="p">):</span>
+    <span class="sd">&quot;&quot;&quot;</span>
+<span class="sd">    Raise a clear error if X is a pandas dataframe</span>
+<span class="sd">    and check array according to scikit rules</span>
+<span class="sd">    &quot;&quot;&quot;</span>
+    <span class="k">if</span> <span class="nb">isinstance</span><span class="p">(</span><span class="n">X</span><span class="p">,</span> <span class="p">(</span><span class="n">pd</span><span class="o">.</span><span class="n">DataFrame</span><span class="p">,</span> <span class="n">pd</span><span class="o">.</span><span class="n">Series</span><span class="p">)):</span>
+        <span class="n">err_message</span> <span class="o">=</span> <span class="s2">&quot;Pandas DataFrame are not supported: apply X.values when calling fit&quot;</span>
+        <span class="k">raise</span> <span class="ne">TypeError</span><span class="p">(</span><span class="n">err_message</span><span class="p">)</span>
+    <span class="n">check_array</span><span class="p">(</span><span class="n">X</span><span class="p">)</span></div>
+
+
+<div class="viewcode-block" id="check_warm_start"><a class="viewcode-back" href="../../generated_docs/pytorch_tabnet.html#pytorch_tabnet.utils.check_warm_start">[docs]</a><span class="k">def</span> <span class="nf">check_warm_start</span><span class="p">(</span><span class="n">warm_start</span><span class="p">,</span> <span class="n">from_unsupervised</span><span class="p">):</span>
+    <span class="sd">&quot;&quot;&quot;</span>
+<span class="sd">    Gives a warning about ambiguous usage of the two parameters.</span>
+<span class="sd">    &quot;&quot;&quot;</span>
+    <span class="k">if</span> <span class="n">warm_start</span> <span class="ow">and</span> <span class="n">from_unsupervised</span> <span class="ow">is</span> <span class="ow">not</span> <span class="kc">None</span><span class="p">:</span>
+        <span class="n">warn_msg</span> <span class="o">=</span> <span class="s2">&quot;warm_start=True and from_unsupervised != None: &quot;</span>
+        <span class="n">warn_msg</span> <span class="o">=</span> <span class="s2">&quot;warm_start will be ignore, training will start from unsupervised weights&quot;</span>
+        <span class="n">warnings</span><span class="o">.</span><span class="n">warn</span><span class="p">(</span><span class="n">warn_msg</span><span class="p">)</span>
+    <span class="k">return</span></div>
 </pre></div>
 
            </div>
diff --git a/docs/_modules/torch/optim/adam.html b/docs/_modules/torch/optim/adam.html
index 2f5a4891..7871f583 100644
--- a/docs/_modules/torch/optim/adam.html
+++ b/docs/_modules/torch/optim/adam.html
@@ -82,14 +82,17 @@
               
             
             
-              <p class="caption"><span class="caption-text">Contents:</span></p>
+              <p><span class="caption-text">Contents:</span></p>
 <ul>
 <li class="toctree-l1"><a class="reference internal" href="../../../generated_docs/README.html">README</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../../generated_docs/README.html#tabnet-attentive-interpretable-tabular-learning">TabNet : Attentive Interpretable Tabular Learning</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../../generated_docs/README.html#installation">Installation</a></li>
-<li class="toctree-l1"><a class="reference internal" href="../../../generated_docs/README.html#what-problems-does-pytorch-tabnet-handles">What problems does pytorch-tabnet handles?</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../../generated_docs/README.html#contributing">Contributing</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../../generated_docs/README.html#what-problems-does-pytorch-tabnet-handle">What problems does pytorch-tabnet handle?</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../../generated_docs/README.html#how-to-use-it">How to use it?</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../../generated_docs/README.html#semi-supervised-pre-training">Semi-supervised pre-training</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../../generated_docs/README.html#data-augmentation-on-the-fly">Data augmentation on the fly</a></li>
+<li class="toctree-l1"><a class="reference internal" href="../../../generated_docs/README.html#easy-saving-and-loading">Easy saving and loading</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../../generated_docs/README.html#useful-links">Useful links</a></li>
 <li class="toctree-l1"><a class="reference internal" href="../../../generated_docs/pytorch_tabnet.html">pytorch_tabnet package</a></li>
 </ul>
diff --git a/docs/_sources/generated_docs/README.md.txt b/docs/_sources/generated_docs/README.md.txt
index 049c9c49..0f0954bb 100644
--- a/docs/_sources/generated_docs/README.md.txt
+++ b/docs/_sources/generated_docs/README.md.txt
@@ -4,21 +4,48 @@
 
 This is a pyTorch implementation of Tabnet (Arik, S. O., & Pfister, T. (2019). TabNet: Attentive Interpretable Tabular Learning. arXiv preprint arXiv:1908.07442.) https://arxiv.org/pdf/1908.07442.pdf.
 
+<!--- BADGES: START --->
 [![CircleCI](https://circleci.com/gh/dreamquark-ai/tabnet.svg?style=svg)](https://circleci.com/gh/dreamquark-ai/tabnet)
 
 [![PyPI version](https://badge.fury.io/py/pytorch-tabnet.svg)](https://badge.fury.io/py/pytorch-tabnet)
 
 ![PyPI - Downloads](https://img.shields.io/pypi/dm/pytorch-tabnet)
 
+[![PyPI - Python Version](https://img.shields.io/pypi/pyversions/pytorch-tabnet?logo=pypi&style=flat&color=blue)][#pypi-package]
+
+[![Conda - Platform](https://img.shields.io/conda/pn/conda-forge/pytorch-tabnet?logo=anaconda&style=flat)][#conda-forge-package]
+
+[![Conda (channel only)](https://img.shields.io/conda/vn/conda-forge/pytorch-tabnet?logo=anaconda&style=flat&color=orange)][#conda-forge-package]
+
+[![GitHub - License](https://img.shields.io/github/license/dreamquark-ai/tabnet?logo=github&style=flat&color=green)][#github-license]
+
+[#github-license]: https://github.com/dreamquark-ai/tabnet/blob/main/LICENSE
+[#pypi-package]: https://pypi.org/project/pytorch-tabnet/
+[#conda-forge-package]: https://anaconda.org/conda-forge/pytorch-tabnet
+<!--- BADGES: END --->
+
 Any questions ? Want to contribute ? To talk with us ? You can join us on [Slack](https://join.slack.com/t/mltooling/shared_invite/zt-fxaj0qk7-SWy2_~EWyhj4x9SD6gbRvg)
 
 # Installation
 
 ## Easy installation
-You can install using pip by running:
-`pip install pytorch-tabnet`
+
+You can install using `pip` or `conda` as follows.
+
+**with pip**
+
+```sh
+pip install pytorch-tabnet
+```
+
+**with conda**
+
+```sh
+conda install -c conda-forge pytorch-tabnet
+```
 
 ## Source code
+
 If you wan to use it locally within a docker container:
 
 - `git clone git@github.com:dreamquark-ai/tabnet.git`
@@ -26,18 +53,28 @@ If you wan to use it locally within a docker container:
 - `cd tabnet` to get inside the repository
 
 -----------------
+
 #### CPU only
+
 - `make start` to build and get inside the container
 
 #### GPU
+
 - `make start-gpu` to build and get inside the GPU container
 
 -----------------
+
 - `poetry install` to install all the dependencies, including jupyter
 
 - `make notebook` inside the same terminal. You can then follow the link to a jupyter notebook with tabnet installed.
 
-# What problems does pytorch-tabnet handles?
+# Contributing
+
+When contributing to the TabNet repository, please make sure to first discuss the change you wish to make via a new or already existing issue.
+
+Our commits follow the rules presented [here](https://www.conventionalcommits.org/en/v1.0.0/).
+
+# What problems does pytorch-tabnet handle?
 
 - TabNetClassifier : binary classification and multi-class classification problems
 - TabNetRegressor : simple and multi-task regression problems
@@ -70,11 +107,11 @@ clf.fit(
 preds = clf.predict(X_test)
 ```
 
-The targets on `y_train/y_valid` should contain a unique type (i.e. they must all be strings or integers).
+The targets on `y_train/y_valid` should contain a unique type (e.g. they must all be strings or integers).
 
 ### Default eval_metric
 
-A few classical evaluation metrics are implemented (see bellow section for custom ones):
+A few classic evaluation metrics are implemented (see further below for custom ones):
 - binary classification metrics : 'auc', 'accuracy', 'balanced_accuracy', 'logloss'
 - multiclass classification : 'accuracy', 'balanced_accuracy', 'logloss'
 - regression: 'mse', 'mae', 'rmse', 'rmsle'
@@ -86,7 +123,7 @@ In order to match the given scores, you need to use `np.clip(clf.predict(X_predi
 
 ### Custom evaluation metrics
 
-It's easy to create a metric that matches your specific need. Here is an example for gini score (note that you need to specifiy whether this metric should be maximized or not):
+You can create a metric for your specific need. Here is an example for gini score (note that you need to specifiy whether this metric should be maximized or not):
 
 ```python
 from pytorch_tabnet.metrics import Metric
@@ -148,13 +185,31 @@ clf.fit(
 )
 ```
 
-The loss function has been normalized to be independent of `pretraining_ratio`, `batch_size` and number of features in the problem.
+The loss function has been normalized to be independent of `pretraining_ratio`, `batch_size` and the number of features in the problem.
 A self supervised loss greater than 1 means that your model is reconstructing worse than predicting the mean for each feature, a loss bellow 1 means that the model is doing better than predicting the mean.
 
 A complete example can be found within the notebook `pretraining_example.ipynb`.
 
 /!\ : current implementation is trying to reconstruct the original inputs, but Batch Normalization applies a random transformation that can't be deduced by a single line, making the reconstruction harder. Lowering the `batch_size` might make the pretraining easier.
 
+# Data augmentation on the fly
+
+It is now possible to apply custom data augmentation pipeline during training.
+Templates for ClassificationSMOTE and RegressionSMOTE have been added in `pytorch-tabnet/augmentations.py` and can be used as is.
+
+
+# Easy saving and loading
+
+It's really easy to save and re-load a trained model, this makes TabNet production ready.
+```
+# save tabnet model
+saving_path_name = "./tabnet_model_test_1"
+saved_filepath = clf.save_model(saving_path_name)
+
+# define new model with basic parameters and load state dict weights
+loaded_clf = TabNetClassifier()
+loaded_clf.load_model(saved_filepath)
+```
 
 # Useful links
 
@@ -237,7 +292,7 @@ A complete example can be found within the notebook `pretraining_example.ipynb`.
 
 - `optimizer_params`: dict (default=dict(lr=2e-2))
 
-    Parameters compatible with optimizer_fn used initialize the optimizer. Since we have Adam as our default optimizer, we use this to define the initial learning rate used for training. As mentionned in the original paper, a large initial learning of ```0.02 ```  with decay is a good option.
+    Parameters compatible with optimizer_fn used initialize the optimizer. Since we have Adam as our default optimizer, we use this to define the initial learning rate used for training. As mentionned in the original paper, a large initial learning rate of ```0.02 ```  with decay is a good option.
 
 - `scheduler_fn` : torch.optim.lr_scheduler (default=None)
 
@@ -251,10 +306,6 @@ A complete example can be found within the notebook `pretraining_example.ipynb`.
 
     Name of the model used for saving in disk, you can customize this to easily retrieve and reuse your trained models.
 
-- `saving_path` : str (default = './')
-
-    Path defining where to save models.
-
 - `verbose` : int (default=1)
 
     Verbosity for notebooks plots, set to 1 to see every epoch, 0 to get None.
@@ -263,7 +314,15 @@ A complete example can be found within the notebook `pretraining_example.ipynb`.
     'cpu' for cpu training, 'gpu' for gpu training, 'auto' to automatically detect gpu.
 
 - `mask_type: str` (default='sparsemax')
-    Either "sparsemax" or "entmax" : this is the masking function to use for selecting features
+    Either "sparsemax" or "entmax" : this is the masking function to use for selecting features.
+
+- `n_shared_decoder` : int (default=1)
+
+    Number of shared GLU block in decoder, this is only useful for `TabNetPretrainer`.
+
+- `n_indep_decoder` : int (default=1)
+
+    Number of independent GLU block in decoder, this is only useful for `TabNetPretrainer`.
 
 ## Fit parameters
 
@@ -291,13 +350,13 @@ A complete example can be found within the notebook `pretraining_example.ipynb`.
 
     Maximum number of epochs for trainng.
     
-- `patience` : int (default = 15)
+- `patience` : int (default = 10)
 
     Number of consecutive epochs without improvement before performing early stopping.
 
-    If patience is set to 0 then no early stopping will be performed.
+    If patience is set to 0, then no early stopping will be performed.
 
-    Note that if patience is enabled, best weights from best epoch will automatically be loaded at the end of `fit`.
+    Note that if patience is enabled, then best weights from best epoch will automatically be loaded at the end of `fit`.
 
 - `weights` : int or dict (default=0)
 
@@ -315,7 +374,7 @@ A complete example can be found within the notebook `pretraining_example.ipynb`.
 
 - `batch_size` : int (default=1024)
 
-    Number of examples per batch, large batch sizes are recommended.
+    Number of examples per batch. Large batch sizes are recommended.
 
 - `virtual_batch_size` : int (default=128)
 
@@ -338,3 +397,7 @@ A complete example can be found within the notebook `pretraining_example.ipynb`.
         /!\ TabNetPretrainer Only : Percentage of input features to mask during pretraining.
 
         Should be between 0 and 1. The bigger the harder the reconstruction task is.
+
+- `warm_start` : bool (default=False)
+    In order to match scikit-learn API, this is set to False.
+    It allows to fit twice the same model and start from a warm start.
diff --git a/docs/_sources/generated_docs/pytorch_tabnet.rst.txt b/docs/_sources/generated_docs/pytorch_tabnet.rst.txt
index 9eca1050..18680c50 100644
--- a/docs/_sources/generated_docs/pytorch_tabnet.rst.txt
+++ b/docs/_sources/generated_docs/pytorch_tabnet.rst.txt
@@ -2,19 +2,19 @@ pytorch_tabnet package
 ======================
 
 
-pytorch_tabnet.utils module
-.............................
+pytorch_tabnet.pretraining_utils module
+.........................................
 
-.. automodule:: pytorch_tabnet.utils
+.. automodule:: pytorch_tabnet.pretraining_utils
     :members:
     :undoc-members:
     :show-inheritance:
 
 
-pytorch_tabnet.pretraining_utils module
-.........................................
+pytorch_tabnet.augmentations module
+.....................................
 
-.. automodule:: pytorch_tabnet.pretraining_utils
+.. automodule:: pytorch_tabnet.augmentations
     :members:
     :undoc-members:
     :show-inheritance:
@@ -29,28 +29,37 @@ pytorch_tabnet.tab_network module
     :show-inheritance:
 
 
-pytorch_tabnet.multitask module
+pytorch_tabnet.metrics module
+...............................
+
+.. automodule:: pytorch_tabnet.metrics
+    :members:
+    :undoc-members:
+    :show-inheritance:
+
+
+pytorch_tabnet.tab_model module
 .................................
 
-.. automodule:: pytorch_tabnet.multitask
+.. automodule:: pytorch_tabnet.tab_model
     :members:
     :undoc-members:
     :show-inheritance:
 
 
-pytorch_tabnet.multiclass_utils module
-........................................
+pytorch_tabnet.sparsemax module
+.................................
 
-.. automodule:: pytorch_tabnet.multiclass_utils
+.. automodule:: pytorch_tabnet.sparsemax
     :members:
     :undoc-members:
     :show-inheritance:
 
 
-pytorch_tabnet.tab_model module
+pytorch_tabnet.callbacks module
 .................................
 
-.. automodule:: pytorch_tabnet.tab_model
+.. automodule:: pytorch_tabnet.callbacks
     :members:
     :undoc-members:
     :show-inheritance:
@@ -65,37 +74,37 @@ pytorch_tabnet.abstract_model module
     :show-inheritance:
 
 
-pytorch_tabnet.callbacks module
-.................................
+pytorch_tabnet.pretraining module
+...................................
 
-.. automodule:: pytorch_tabnet.callbacks
+.. automodule:: pytorch_tabnet.pretraining
     :members:
     :undoc-members:
     :show-inheritance:
 
 
-pytorch_tabnet.pretraining module
-...................................
+pytorch_tabnet.utils module
+.............................
 
-.. automodule:: pytorch_tabnet.pretraining
+.. automodule:: pytorch_tabnet.utils
     :members:
     :undoc-members:
     :show-inheritance:
 
 
-pytorch_tabnet.sparsemax module
+pytorch_tabnet.multitask module
 .................................
 
-.. automodule:: pytorch_tabnet.sparsemax
+.. automodule:: pytorch_tabnet.multitask
     :members:
     :undoc-members:
     :show-inheritance:
 
 
-pytorch_tabnet.metrics module
-...............................
+pytorch_tabnet.multiclass_utils module
+........................................
 
-.. automodule:: pytorch_tabnet.metrics
+.. automodule:: pytorch_tabnet.multiclass_utils
     :members:
     :undoc-members:
     :show-inheritance:
diff --git a/docs/_static/pygments.css b/docs/_static/pygments.css
index d14395ef..691aeb82 100644
--- a/docs/_static/pygments.css
+++ b/docs/_static/pygments.css
@@ -1,8 +1,8 @@
-pre { line-height: 125%; margin: 0; }
-td.linenos pre { color: #000000; background-color: #f0f0f0; padding: 0 5px 0 5px; }
-span.linenos { color: #000000; background-color: #f0f0f0; padding: 0 5px 0 5px; }
-td.linenos pre.special { color: #000000; background-color: #ffffc0; padding: 0 5px 0 5px; }
-span.linenos.special { color: #000000; background-color: #ffffc0; padding: 0 5px 0 5px; }
+pre { line-height: 125%; }
+td.linenos .normal { color: inherit; background-color: transparent; padding-left: 5px; padding-right: 5px; }
+span.linenos { color: inherit; background-color: transparent; padding-left: 5px; padding-right: 5px; }
+td.linenos .special { color: #000000; background-color: #ffffc0; padding-left: 5px; padding-right: 5px; }
+span.linenos.special { color: #000000; background-color: #ffffc0; padding-left: 5px; padding-right: 5px; }
 .highlight .hll { background-color: #ffffcc }
 .highlight { background: #eeffcc; }
 .highlight .c { color: #408090; font-style: italic } /* Comment */
diff --git a/docs/generated_docs/README.html b/docs/generated_docs/README.html
index 329f3a77..f6deca28 100644
--- a/docs/generated_docs/README.html
+++ b/docs/generated_docs/README.html
@@ -4,7 +4,8 @@
 <html class="writer-html5" lang="en" >
 <head>
   <meta charset="utf-8">
-  
+  <meta name="generator" content="Docutils 0.17.1: http://docutils.sourceforge.net/" />
+
   <meta name="viewport" content="width=device-width, initial-scale=1.0">
   
   <title>README &mdash; pytorch_tabnet  documentation</title>
@@ -84,7 +85,7 @@
               
             
             
-              <p class="caption"><span class="caption-text">Contents:</span></p>
+              <p><span class="caption-text">Contents:</span></p>
 <ul class="current">
 <li class="toctree-l1 current"><a class="current reference internal" href="#">README</a></li>
 <li class="toctree-l1"><a class="reference internal" href="#tabnet-attentive-interpretable-tabular-learning">TabNet : Attentive Interpretable Tabular Learning</a></li>
@@ -97,13 +98,16 @@
 </li>
 </ul>
 </li>
-<li class="toctree-l1"><a class="reference internal" href="#what-problems-does-pytorch-tabnet-handles">What problems does pytorch-tabnet handles?</a></li>
+<li class="toctree-l1"><a class="reference internal" href="#contributing">Contributing</a></li>
+<li class="toctree-l1"><a class="reference internal" href="#what-problems-does-pytorch-tabnet-handle">What problems does pytorch-tabnet handle?</a></li>
 <li class="toctree-l1"><a class="reference internal" href="#how-to-use-it">How to use it?</a><ul>
 <li class="toctree-l2"><a class="reference internal" href="#default-eval-metric">Default eval_metric</a></li>
 <li class="toctree-l2"><a class="reference internal" href="#custom-evaluation-metrics">Custom evaluation metrics</a></li>
 </ul>
 </li>
 <li class="toctree-l1"><a class="reference internal" href="#semi-supervised-pre-training">Semi-supervised pre-training</a></li>
+<li class="toctree-l1"><a class="reference internal" href="#data-augmentation-on-the-fly">Data augmentation on the fly</a></li>
+<li class="toctree-l1"><a class="reference internal" href="#easy-saving-and-loading">Easy saving and loading</a></li>
 <li class="toctree-l1"><a class="reference internal" href="#useful-links">Useful links</a><ul>
 <li class="toctree-l2"><a class="reference internal" href="#model-parameters">Model parameters</a></li>
 <li class="toctree-l2"><a class="reference internal" href="#fit-parameters">Fit parameters</a></li>
@@ -175,25 +179,36 @@
           <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
            <div itemprop="articleBody">
             
-  <div class="section" id="readme">
+  <section id="readme">
 <h1>README<a class="headerlink" href="#readme" title="Permalink to this headline">¶</a></h1>
-</div>
-<div class="section" id="tabnet-attentive-interpretable-tabular-learning">
+</section>
+<section id="tabnet-attentive-interpretable-tabular-learning">
 <h1>TabNet : Attentive Interpretable Tabular Learning<a class="headerlink" href="#tabnet-attentive-interpretable-tabular-learning" title="Permalink to this headline">¶</a></h1>
 <p>This is a pyTorch implementation of Tabnet (Arik, S. O., &amp; Pfister, T. (2019). TabNet: Attentive Interpretable Tabular Learning. arXiv preprint arXiv:1908.07442.) https://arxiv.org/pdf/1908.07442.pdf.</p>
-<p><a class="reference external" href="https://circleci.com/gh/dreamquark-ai/tabnet"><img alt="CircleCI" src="https://circleci.com/gh/dreamquark-ai/tabnet.svg?style=svg" /></a></p>
+<!--- BADGES: START ---><p><a class="reference external" href="https://circleci.com/gh/dreamquark-ai/tabnet"><img alt="CircleCI" src="https://circleci.com/gh/dreamquark-ai/tabnet.svg?style=svg" /></a></p>
 <p><a class="reference external" href="https://badge.fury.io/py/pytorch-tabnet"><img alt="PyPI version" src="https://badge.fury.io/py/pytorch-tabnet.svg" /></a></p>
 <p><img alt="PyPI - Downloads" src="https://img.shields.io/pypi/dm/pytorch-tabnet" /></p>
-<p>Any questions ? Want to contribute ? To talk with us ? You can join us on <a class="reference external" href="https://join.slack.com/t/mltooling/shared_invite/zt-fxaj0qk7-SWy2_~EWyhj4x9SD6gbRvg">Slack</a></p>
-</div>
-<div class="section" id="installation">
+<p><a class="reference external" href="https://pypi.org/project/pytorch-tabnet/"><img alt="PyPI - Python Version" src="https://img.shields.io/pypi/pyversions/pytorch-tabnet?logo=pypi&amp;style=flat&amp;color=blue" /></a></p>
+<p><a class="reference external" href="https://anaconda.org/conda-forge/pytorch-tabnet"><img alt="Conda - Platform" src="https://img.shields.io/conda/pn/conda-forge/pytorch-tabnet?logo=anaconda&amp;style=flat" /></a></p>
+<p><a class="reference external" href="https://anaconda.org/conda-forge/pytorch-tabnet"><img alt="Conda (channel only)" src="https://img.shields.io/conda/vn/conda-forge/pytorch-tabnet?logo=anaconda&amp;style=flat&amp;color=orange" /></a></p>
+<p><a class="reference external" href="https://github.com/dreamquark-ai/tabnet/blob/main/LICENSE"><img alt="GitHub - License" src="https://img.shields.io/github/license/dreamquark-ai/tabnet?logo=github&amp;style=flat&amp;color=green" /></a></p>
+<!--- BADGES: END ---><p>Any questions ? Want to contribute ? To talk with us ? You can join us on <a class="reference external" href="https://join.slack.com/t/mltooling/shared_invite/zt-fxaj0qk7-SWy2_~EWyhj4x9SD6gbRvg">Slack</a></p>
+</section>
+<section id="installation">
 <h1>Installation<a class="headerlink" href="#installation" title="Permalink to this headline">¶</a></h1>
-<div class="section" id="easy-installation">
+<section id="easy-installation">
 <h2>Easy installation<a class="headerlink" href="#easy-installation" title="Permalink to this headline">¶</a></h2>
-<p>You can install using pip by running:
-<code class="docutils literal notranslate"><span class="pre">pip</span> <span class="pre">install</span> <span class="pre">pytorch-tabnet</span></code></p>
+<p>You can install using <code class="docutils literal notranslate"><span class="pre">pip</span></code> or <code class="docutils literal notranslate"><span class="pre">conda</span></code> as follows.</p>
+<p><strong>with pip</strong></p>
+<div class="highlight-sh notranslate"><div class="highlight"><pre><span></span>pip install pytorch-tabnet
+</pre></div>
 </div>
-<div class="section" id="source-code">
+<p><strong>with conda</strong></p>
+<div class="highlight-sh notranslate"><div class="highlight"><pre><span></span>conda install -c conda-forge pytorch-tabnet
+</pre></div>
+</div>
+</section>
+<section id="source-code">
 <h2>Source code<a class="headerlink" href="#source-code" title="Permalink to this headline">¶</a></h2>
 <p>If you wan to use it locally within a docker container:</p>
 <ul class="simple">
@@ -201,13 +216,13 @@ <h2>Source code<a class="headerlink" href="#source-code" title="Permalink to thi
 <li><p><code class="docutils literal notranslate"><span class="pre">cd</span> <span class="pre">tabnet</span></code> to get inside the repository</p></li>
 </ul>
 <hr class="docutils" />
-<div class="section" id="cpu-only">
+<section id="cpu-only">
 <h3>CPU only<a class="headerlink" href="#cpu-only" title="Permalink to this headline">¶</a></h3>
 <ul class="simple">
 <li><p><code class="docutils literal notranslate"><span class="pre">make</span> <span class="pre">start</span></code> to build and get inside the container</p></li>
 </ul>
-</div>
-<div class="section" id="gpu">
+</section>
+<section id="gpu">
 <h3>GPU<a class="headerlink" href="#gpu" title="Permalink to this headline">¶</a></h3>
 <ul class="simple">
 <li><p><code class="docutils literal notranslate"><span class="pre">make</span> <span class="pre">start-gpu</span></code> to build and get inside the GPU container</p></li>
@@ -217,18 +232,23 @@ <h3>GPU<a class="headerlink" href="#gpu" title="Permalink to this headline">¶</
 <li><p><code class="docutils literal notranslate"><span class="pre">poetry</span> <span class="pre">install</span></code> to install all the dependencies, including jupyter</p></li>
 <li><p><code class="docutils literal notranslate"><span class="pre">make</span> <span class="pre">notebook</span></code> inside the same terminal. You can then follow the link to a jupyter notebook with tabnet installed.</p></li>
 </ul>
-</div>
-</div>
-</div>
-<div class="section" id="what-problems-does-pytorch-tabnet-handles">
-<h1>What problems does pytorch-tabnet handles?<a class="headerlink" href="#what-problems-does-pytorch-tabnet-handles" title="Permalink to this headline">¶</a></h1>
+</section>
+</section>
+</section>
+<section id="contributing">
+<h1>Contributing<a class="headerlink" href="#contributing" title="Permalink to this headline">¶</a></h1>
+<p>When contributing to the TabNet repository, please make sure to first discuss the change you wish to make via a new or already existing issue.</p>
+<p>Our commits follow the rules presented <a class="reference external" href="https://www.conventionalcommits.org/en/v1.0.0/">here</a>.</p>
+</section>
+<section id="what-problems-does-pytorch-tabnet-handle">
+<h1>What problems does pytorch-tabnet handle?<a class="headerlink" href="#what-problems-does-pytorch-tabnet-handle" title="Permalink to this headline">¶</a></h1>
 <ul class="simple">
 <li><p>TabNetClassifier : binary classification and multi-class classification problems</p></li>
 <li><p>TabNetRegressor : simple and multi-task regression problems</p></li>
 <li><p>TabNetMultiTaskClassifier:  multi-task multi-classification problems</p></li>
 </ul>
-</div>
-<div class="section" id="how-to-use-it">
+</section>
+<section id="how-to-use-it">
 <h1>How to use it?<a class="headerlink" href="#how-to-use-it" title="Permalink to this headline">¶</a></h1>
 <p>TabNet is now scikit-compatible, training a TabNetClassifier or TabNetRegressor is really easy.</p>
 <div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="kn">from</span> <span class="nn">pytorch_tabnet.tab_model</span> <span class="kn">import</span> <span class="n">TabNetClassifier</span><span class="p">,</span> <span class="n">TabNetRegressor</span>
@@ -251,10 +271,10 @@ <h1>How to use it?<a class="headerlink" href="#how-to-use-it" title="Permalink t
 <span class="n">preds</span> <span class="o">=</span> <span class="n">clf</span><span class="o">.</span><span class="n">predict</span><span class="p">(</span><span class="n">X_test</span><span class="p">)</span>
 </pre></div>
 </div>
-<p>The targets on <code class="docutils literal notranslate"><span class="pre">y_train/y_valid</span></code> should contain a unique type (i.e. they must all be strings or integers).</p>
-<div class="section" id="default-eval-metric">
+<p>The targets on <code class="docutils literal notranslate"><span class="pre">y_train/y_valid</span></code> should contain a unique type (e.g. they must all be strings or integers).</p>
+<section id="default-eval-metric">
 <h2>Default eval_metric<a class="headerlink" href="#default-eval-metric" title="Permalink to this headline">¶</a></h2>
-<p>A few classical evaluation metrics are implemented (see bellow section for custom ones):</p>
+<p>A few classic evaluation metrics are implemented (see further below for custom ones):</p>
 <ul class="simple">
 <li><p>binary classification metrics : ‘auc’, ‘accuracy’, ‘balanced_accuracy’, ‘logloss’</p></li>
 <li><p>multiclass classification : ‘accuracy’, ‘balanced_accuracy’, ‘logloss’</p></li>
@@ -262,10 +282,10 @@ <h2>Default eval_metric<a class="headerlink" href="#default-eval-metric" title="
 </ul>
 <p>Important Note : ‘rmsle’ will automatically clip negative predictions to 0, because the model can predict negative values.
 In order to match the given scores, you need to use <code class="docutils literal notranslate"><span class="pre">np.clip(clf.predict(X_predict),</span> <span class="pre">a_min=0,</span> <span class="pre">a_max=None)</span></code> when doing predictions.</p>
-</div>
-<div class="section" id="custom-evaluation-metrics">
+</section>
+<section id="custom-evaluation-metrics">
 <h2>Custom evaluation metrics<a class="headerlink" href="#custom-evaluation-metrics" title="Permalink to this headline">¶</a></h2>
-<p>It’s easy to create a metric that matches your specific need. Here is an example for gini score (note that you need to specifiy whether this metric should be maximized or not):</p>
+<p>You can create a metric for your specific need. Here is an example for gini score (note that you need to specifiy whether this metric should be maximized or not):</p>
 <div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="kn">from</span> <span class="nn">pytorch_tabnet.metrics</span> <span class="kn">import</span> <span class="n">Metric</span>
 <span class="kn">from</span> <span class="nn">sklearn.metrics</span> <span class="kn">import</span> <span class="n">roc_auc_score</span>
 
@@ -287,9 +307,9 @@ <h2>Custom evaluation metrics<a class="headerlink" href="#custom-evaluation-metr
 </pre></div>
 </div>
 <p>A specific customization example notebook is available here : https://github.com/dreamquark-ai/tabnet/blob/develop/customizing_example.ipynb</p>
-</div>
-</div>
-<div class="section" id="semi-supervised-pre-training">
+</section>
+</section>
+<section id="semi-supervised-pre-training">
 <h1>Semi-supervised pre-training<a class="headerlink" href="#semi-supervised-pre-training" title="Permalink to this headline">¶</a></h1>
 <p>Added later to TabNet’s original paper, semi-supervised pre-training is now available via the class <code class="docutils literal notranslate"><span class="pre">TabNetPretrainer</span></code>:</p>
 <div class="highlight-python notranslate"><div class="highlight"><pre><span></span><span class="c1"># TabNetPretrainer</span>
@@ -323,12 +343,30 @@ <h1>Semi-supervised pre-training<a class="headerlink" href="#semi-supervised-pre
 <span class="p">)</span>
 </pre></div>
 </div>
-<p>The loss function has been normalized to be independent of <code class="docutils literal notranslate"><span class="pre">pretraining_ratio</span></code>, <code class="docutils literal notranslate"><span class="pre">batch_size</span></code> and number of features in the problem.
+<p>The loss function has been normalized to be independent of <code class="docutils literal notranslate"><span class="pre">pretraining_ratio</span></code>, <code class="docutils literal notranslate"><span class="pre">batch_size</span></code> and the number of features in the problem.
 A self supervised loss greater than 1 means that your model is reconstructing worse than predicting the mean for each feature, a loss bellow 1 means that the model is doing better than predicting the mean.</p>
 <p>A complete example can be found within the notebook <code class="docutils literal notranslate"><span class="pre">pretraining_example.ipynb</span></code>.</p>
 <p>/!\ : current implementation is trying to reconstruct the original inputs, but Batch Normalization applies a random transformation that can’t be deduced by a single line, making the reconstruction harder. Lowering the <code class="docutils literal notranslate"><span class="pre">batch_size</span></code> might make the pretraining easier.</p>
+</section>
+<section id="data-augmentation-on-the-fly">
+<h1>Data augmentation on the fly<a class="headerlink" href="#data-augmentation-on-the-fly" title="Permalink to this headline">¶</a></h1>
+<p>It is now possible to apply custom data augmentation pipeline during training.
+Templates for ClassificationSMOTE and RegressionSMOTE have been added in <code class="docutils literal notranslate"><span class="pre">pytorch-tabnet/augmentations.py</span></code> and can be used as is.</p>
+</section>
+<section id="easy-saving-and-loading">
+<h1>Easy saving and loading<a class="headerlink" href="#easy-saving-and-loading" title="Permalink to this headline">¶</a></h1>
+<p>It’s really easy to save and re-load a trained model, this makes TabNet production ready.</p>
+<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="c1"># save tabnet model</span>
+<span class="n">saving_path_name</span> <span class="o">=</span> <span class="s2">&quot;./tabnet_model_test_1&quot;</span>
+<span class="n">saved_filepath</span> <span class="o">=</span> <span class="n">clf</span><span class="o">.</span><span class="n">save_model</span><span class="p">(</span><span class="n">saving_path_name</span><span class="p">)</span>
+
+<span class="c1"># define new model with basic parameters and load state dict weights</span>
+<span class="n">loaded_clf</span> <span class="o">=</span> <span class="n">TabNetClassifier</span><span class="p">()</span>
+<span class="n">loaded_clf</span><span class="o">.</span><span class="n">load_model</span><span class="p">(</span><span class="n">saved_filepath</span><span class="p">)</span>
+</pre></div>
 </div>
-<div class="section" id="useful-links">
+</section>
+<section id="useful-links">
 <h1>Useful links<a class="headerlink" href="#useful-links" title="Permalink to this headline">¶</a></h1>
 <ul class="simple">
 <li><p><a class="reference external" href="https://youtu.be/ysBaZO8YmX8">explanatory video</a></p></li>
@@ -339,7 +377,7 @@ <h1>Useful links<a class="headerlink" href="#useful-links" title="Permalink to t
 <li><p><a class="reference external" href="https://www.kaggle.com/optimo/tabnetmultitaskclassifier">multi-task multi-class classification examples</a></p></li>
 <li><p><a class="reference external" href="https://www.kaggle.com/c/lish-moa/discussion/201510">kaggle moa 1st place solution using tabnet</a></p></li>
 </ul>
-<div class="section" id="model-parameters">
+<section id="model-parameters">
 <h2>Model parameters<a class="headerlink" href="#model-parameters" title="Permalink to this headline">¶</a></h2>
 <ul>
 <li><p><code class="docutils literal notranslate"><span class="pre">n_d</span></code> : int (default=8)</p>
@@ -395,30 +433,33 @@ <h2>Model parameters<a class="headerlink" href="#model-parameters" title="Permal
 <p>Pytorch optimizer function</p>
 </li>
 <li><p><code class="docutils literal notranslate"><span class="pre">optimizer_params</span></code>: dict (default=dict(lr=2e-2))</p>
-<p>Parameters compatible with optimizer_fn used initialize the optimizer. Since we have Adam as our default optimizer, we use this to define the initial learning rate used for training. As mentionned in the original paper, a large initial learning of <code class="docutils literal notranslate"><span class="pre">0.02</span> </code>  with decay is a good option.</p>
+<p>Parameters compatible with optimizer_fn used initialize the optimizer. Since we have Adam as our default optimizer, we use this to define the initial learning rate used for training. As mentionned in the original paper, a large initial learning rate of <code class="docutils literal notranslate"><span class="pre">0.02</span> </code>  with decay is a good option.</p>
 </li>
 <li><p><code class="docutils literal notranslate"><span class="pre">scheduler_fn</span></code> : torch.optim.lr_scheduler (default=None)</p>
 <p>Pytorch Scheduler to change learning rates during training.</p>
 </li>
 <li><p><code class="docutils literal notranslate"><span class="pre">scheduler_params</span></code> : dict</p>
-<p>Dictionnary of parameters to apply to the scheduler_fn. Ex : {“gamma”: 0.95, “step_size”: 10}</p>
+<p>Dictionnary of parameters to apply to the scheduler_fn. Ex : {”gamma”: 0.95, “step_size”: 10}</p>
 </li>
 <li><p><code class="docutils literal notranslate"><span class="pre">model_name</span></code> : str (default = ‘DreamQuarkTabNet’)</p>
 <p>Name of the model used for saving in disk, you can customize this to easily retrieve and reuse your trained models.</p>
 </li>
-<li><p><code class="docutils literal notranslate"><span class="pre">saving_path</span></code> : str (default = ‘./’)</p>
-<p>Path defining where to save models.</p>
-</li>
 <li><p><code class="docutils literal notranslate"><span class="pre">verbose</span></code> : int (default=1)</p>
 <p>Verbosity for notebooks plots, set to 1 to see every epoch, 0 to get None.</p>
 </li>
 <li><p><code class="docutils literal notranslate"><span class="pre">device_name</span></code> : str (default=’auto’)
 ‘cpu’ for cpu training, ‘gpu’ for gpu training, ‘auto’ to automatically detect gpu.</p></li>
 <li><p><code class="docutils literal notranslate"><span class="pre">mask_type:</span> <span class="pre">str</span></code> (default=’sparsemax’)
-Either “sparsemax” or “entmax” : this is the masking function to use for selecting features</p></li>
+Either “sparsemax” or “entmax” : this is the masking function to use for selecting features.</p></li>
+<li><p><code class="docutils literal notranslate"><span class="pre">n_shared_decoder</span></code> : int (default=1)</p>
+<p>Number of shared GLU block in decoder, this is only useful for <code class="docutils literal notranslate"><span class="pre">TabNetPretrainer</span></code>.</p>
+</li>
+<li><p><code class="docutils literal notranslate"><span class="pre">n_indep_decoder</span></code> : int (default=1)</p>
+<p>Number of independent GLU block in decoder, this is only useful for <code class="docutils literal notranslate"><span class="pre">TabNetPretrainer</span></code>.</p>
+</li>
 </ul>
-</div>
-<div class="section" id="fit-parameters">
+</section>
+<section id="fit-parameters">
 <h2>Fit parameters<a class="headerlink" href="#fit-parameters" title="Permalink to this headline">¶</a></h2>
 <ul>
 <li><p><code class="docutils literal notranslate"><span class="pre">X_train</span></code> : np.array</p>
@@ -435,10 +476,10 @@ <h2>Fit parameters<a class="headerlink" href="#fit-parameters" title="Permalink
 <li><p><code class="docutils literal notranslate"><span class="pre">max_epochs</span></code> : int (default = 200)</p>
 <p>Maximum number of epochs for trainng.</p>
 </li>
-<li><p><code class="docutils literal notranslate"><span class="pre">patience</span></code> : int (default = 15)</p>
+<li><p><code class="docutils literal notranslate"><span class="pre">patience</span></code> : int (default = 10)</p>
 <p>Number of consecutive epochs without improvement before performing early stopping.</p>
-<p>If patience is set to 0 then no early stopping will be performed.</p>
-<p>Note that if patience is enabled, best weights from best epoch will automatically be loaded at the end of <code class="docutils literal notranslate"><span class="pre">fit</span></code>.</p>
+<p>If patience is set to 0, then no early stopping will be performed.</p>
+<p>Note that if patience is enabled, then best weights from best epoch will automatically be loaded at the end of <code class="docutils literal notranslate"><span class="pre">fit</span></code>.</p>
 </li>
 <li><p><code class="docutils literal notranslate"><span class="pre">weights</span></code> : int or dict (default=0)</p>
 <p>/!\ Only for TabNetClassifier
@@ -453,7 +494,7 @@ <h2>Fit parameters<a class="headerlink" href="#fit-parameters" title="Permalink
 each task will be assigned its own loss function</p>
 </li>
 <li><p><code class="docutils literal notranslate"><span class="pre">batch_size</span></code> : int (default=1024)</p>
-<p>Number of examples per batch, large batch sizes are recommended.</p>
+<p>Number of examples per batch. Large batch sizes are recommended.</p>
 </li>
 <li><p><code class="docutils literal notranslate"><span class="pre">virtual_batch_size</span></code> : int (default=128)</p>
 <p>Size of the mini batches used for “Ghost Batch Normalization”.
@@ -473,9 +514,12 @@ <h2>Fit parameters<a class="headerlink" href="#fit-parameters" title="Permalink
 </pre></div>
 </div>
 </li>
+<li><p><code class="docutils literal notranslate"><span class="pre">warm_start</span></code> : bool (default=False)
+In order to match scikit-learn API, this is set to False.
+It allows to fit twice the same model and start from a warm start.</p></li>
 </ul>
-</div>
-</div>
+</section>
+</section>
 
 
            </div>
diff --git a/docs/generated_docs/pytorch_tabnet.html b/docs/generated_docs/pytorch_tabnet.html
index 5583b3b0..2aefc158 100644
--- a/docs/generated_docs/pytorch_tabnet.html
+++ b/docs/generated_docs/pytorch_tabnet.html
@@ -4,7 +4,8 @@
 <html class="writer-html5" lang="en" >
 <head>
   <meta charset="utf-8">
-  
+  <meta name="generator" content="Docutils 0.17.1: http://docutils.sourceforge.net/" />
+
   <meta name="viewport" content="width=device-width, initial-scale=1.0">
   
   <title>pytorch_tabnet package &mdash; pytorch_tabnet  documentation</title>
@@ -83,30 +84,34 @@
               
             
             
-              <p class="caption"><span class="caption-text">Contents:</span></p>
+              <p><span class="caption-text">Contents:</span></p>
 <ul class="current">
 <li class="toctree-l1"><a class="reference internal" href="README.html">README</a></li>
 <li class="toctree-l1"><a class="reference internal" href="README.html#tabnet-attentive-interpretable-tabular-learning">TabNet : Attentive Interpretable Tabular Learning</a></li>
 <li class="toctree-l1"><a class="reference internal" href="README.html#installation">Installation</a></li>
-<li class="toctree-l1"><a class="reference internal" href="README.html#what-problems-does-pytorch-tabnet-handles">What problems does pytorch-tabnet handles?</a></li>
+<li class="toctree-l1"><a class="reference internal" href="README.html#contributing">Contributing</a></li>
+<li class="toctree-l1"><a class="reference internal" href="README.html#what-problems-does-pytorch-tabnet-handle">What problems does pytorch-tabnet handle?</a></li>
 <li class="toctree-l1"><a class="reference internal" href="README.html#how-to-use-it">How to use it?</a></li>
 <li class="toctree-l1"><a class="reference internal" href="README.html#semi-supervised-pre-training">Semi-supervised pre-training</a></li>
+<li class="toctree-l1"><a class="reference internal" href="README.html#data-augmentation-on-the-fly">Data augmentation on the fly</a></li>
+<li class="toctree-l1"><a class="reference internal" href="README.html#easy-saving-and-loading">Easy saving and loading</a></li>
 <li class="toctree-l1"><a class="reference internal" href="README.html#useful-links">Useful links</a></li>
 <li class="toctree-l1 current"><a class="current reference internal" href="#">pytorch_tabnet package</a><ul>
-<li class="toctree-l2"><a class="reference internal" href="#module-pytorch_tabnet.utils">pytorch_tabnet.utils module</a></li>
 <li class="toctree-l2"><a class="reference internal" href="#module-pytorch_tabnet.pretraining_utils">pytorch_tabnet.pretraining_utils module</a></li>
+<li class="toctree-l2"><a class="reference internal" href="#module-pytorch_tabnet.augmentations">pytorch_tabnet.augmentations module</a></li>
 <li class="toctree-l2"><a class="reference internal" href="#module-pytorch_tabnet.tab_network">pytorch_tabnet.tab_network module</a></li>
+<li class="toctree-l2"><a class="reference internal" href="#module-pytorch_tabnet.metrics">pytorch_tabnet.metrics module</a></li>
+<li class="toctree-l2"><a class="reference internal" href="#module-pytorch_tabnet.tab_model">pytorch_tabnet.tab_model module</a></li>
+<li class="toctree-l2"><a class="reference internal" href="#module-pytorch_tabnet.sparsemax">pytorch_tabnet.sparsemax module</a></li>
+<li class="toctree-l2"><a class="reference internal" href="#module-pytorch_tabnet.callbacks">pytorch_tabnet.callbacks module</a></li>
+<li class="toctree-l2"><a class="reference internal" href="#module-pytorch_tabnet.abstract_model">pytorch_tabnet.abstract_model module</a></li>
+<li class="toctree-l2"><a class="reference internal" href="#module-pytorch_tabnet.pretraining">pytorch_tabnet.pretraining module</a></li>
+<li class="toctree-l2"><a class="reference internal" href="#module-pytorch_tabnet.utils">pytorch_tabnet.utils module</a></li>
 <li class="toctree-l2"><a class="reference internal" href="#module-pytorch_tabnet.multitask">pytorch_tabnet.multitask module</a></li>
 <li class="toctree-l2"><a class="reference internal" href="#module-pytorch_tabnet.multiclass_utils">pytorch_tabnet.multiclass_utils module</a><ul>
 <li class="toctree-l3"><a class="reference internal" href="#multi-class-multi-label-utility-function">Multi-class / multi-label utility function</a></li>
 </ul>
 </li>
-<li class="toctree-l2"><a class="reference internal" href="#module-pytorch_tabnet.tab_model">pytorch_tabnet.tab_model module</a></li>
-<li class="toctree-l2"><a class="reference internal" href="#module-pytorch_tabnet.abstract_model">pytorch_tabnet.abstract_model module</a></li>
-<li class="toctree-l2"><a class="reference internal" href="#module-pytorch_tabnet.callbacks">pytorch_tabnet.callbacks module</a></li>
-<li class="toctree-l2"><a class="reference internal" href="#module-pytorch_tabnet.pretraining">pytorch_tabnet.pretraining module</a></li>
-<li class="toctree-l2"><a class="reference internal" href="#module-pytorch_tabnet.sparsemax">pytorch_tabnet.sparsemax module</a></li>
-<li class="toctree-l2"><a class="reference internal" href="#module-pytorch_tabnet.metrics">pytorch_tabnet.metrics module</a></li>
 </ul>
 </li>
 </ul>
@@ -174,208 +179,9 @@
           <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
            <div itemprop="articleBody">
             
-  <div class="section" id="pytorch-tabnet-package">
+  <section id="pytorch-tabnet-package">
 <h1>pytorch_tabnet package<a class="headerlink" href="#pytorch-tabnet-package" title="Permalink to this headline">¶</a></h1>
-<div class="section" id="module-pytorch_tabnet.utils">
-<span id="pytorch-tabnet-utils-module"></span><h2>pytorch_tabnet.utils module<a class="headerlink" href="#module-pytorch_tabnet.utils" title="Permalink to this headline">¶</a></h2>
-<dl class="class">
-<dt id="pytorch_tabnet.utils.ComplexEncoder">
-<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.utils.</code><code class="sig-name descname">ComplexEncoder</code><span class="sig-paren">(</span><em class="sig-param">*</em>, <em class="sig-param">skipkeys=False</em>, <em class="sig-param">ensure_ascii=True</em>, <em class="sig-param">check_circular=True</em>, <em class="sig-param">allow_nan=True</em>, <em class="sig-param">sort_keys=False</em>, <em class="sig-param">indent=None</em>, <em class="sig-param">separators=None</em>, <em class="sig-param">default=None</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/utils.html#ComplexEncoder"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.utils.ComplexEncoder" title="Permalink to this definition">¶</a></dt>
-<dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">json.encoder.JSONEncoder</span></code></p>
-<dl class="method">
-<dt id="pytorch_tabnet.utils.ComplexEncoder.default">
-<code class="sig-name descname">default</code><span class="sig-paren">(</span><em class="sig-param">obj</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/utils.html#ComplexEncoder.default"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.utils.ComplexEncoder.default" title="Permalink to this definition">¶</a></dt>
-<dd><p>Implement this method in a subclass such that it returns
-a serializable object for <code class="docutils literal notranslate"><span class="pre">o</span></code>, or calls the base implementation
-(to raise a <code class="docutils literal notranslate"><span class="pre">TypeError</span></code>).</p>
-<p>For example, to support arbitrary iterators, you could
-implement default like this:</p>
-<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="k">def</span> <span class="nf">default</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">o</span><span class="p">):</span>
-    <span class="k">try</span><span class="p">:</span>
-        <span class="n">iterable</span> <span class="o">=</span> <span class="nb">iter</span><span class="p">(</span><span class="n">o</span><span class="p">)</span>
-    <span class="k">except</span> <span class="ne">TypeError</span><span class="p">:</span>
-        <span class="k">pass</span>
-    <span class="k">else</span><span class="p">:</span>
-        <span class="k">return</span> <span class="nb">list</span><span class="p">(</span><span class="n">iterable</span><span class="p">)</span>
-    <span class="c1"># Let the base class default method raise the TypeError</span>
-    <span class="k">return</span> <span class="n">JSONEncoder</span><span class="o">.</span><span class="n">default</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">o</span><span class="p">)</span>
-</pre></div>
-</div>
-</dd></dl>
-
-</dd></dl>
-
-<dl class="class">
-<dt id="pytorch_tabnet.utils.PredictDataset">
-<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.utils.</code><code class="sig-name descname">PredictDataset</code><span class="sig-paren">(</span><em class="sig-param">x</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/utils.html#PredictDataset"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.utils.PredictDataset" title="Permalink to this definition">¶</a></dt>
-<dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">torch.utils.data.dataset.Dataset</span></code></p>
-<p>Format for numpy array</p>
-<dl class="field-list simple">
-<dt class="field-odd">Parameters</dt>
-<dd class="field-odd"><p><strong>X</strong> (<em>2D array</em>) – The input matrix</p>
-</dd>
-</dl>
-</dd></dl>
-
-<dl class="class">
-<dt id="pytorch_tabnet.utils.TorchDataset">
-<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.utils.</code><code class="sig-name descname">TorchDataset</code><span class="sig-paren">(</span><em class="sig-param">x</em>, <em class="sig-param">y</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/utils.html#TorchDataset"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.utils.TorchDataset" title="Permalink to this definition">¶</a></dt>
-<dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">torch.utils.data.dataset.Dataset</span></code></p>
-<p>Format for numpy array</p>
-<dl class="field-list simple">
-<dt class="field-odd">Parameters</dt>
-<dd class="field-odd"><ul class="simple">
-<li><p><strong>X</strong> (<em>2D array</em>) – The input matrix</p></li>
-<li><p><strong>y</strong> (<em>2D array</em>) – The one-hot encoded target</p></li>
-</ul>
-</dd>
-</dl>
-</dd></dl>
-
-<dl class="function">
-<dt id="pytorch_tabnet.utils.create_dataloaders">
-<code class="sig-prename descclassname">pytorch_tabnet.utils.</code><code class="sig-name descname">create_dataloaders</code><span class="sig-paren">(</span><em class="sig-param">X_train</em>, <em class="sig-param">y_train</em>, <em class="sig-param">eval_set</em>, <em class="sig-param">weights</em>, <em class="sig-param">batch_size</em>, <em class="sig-param">num_workers</em>, <em class="sig-param">drop_last</em>, <em class="sig-param">pin_memory</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/utils.html#create_dataloaders"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.utils.create_dataloaders" title="Permalink to this definition">¶</a></dt>
-<dd><p>Create dataloaders with or without subsampling depending on weights and balanced.</p>
-<dl class="field-list simple">
-<dt class="field-odd">Parameters</dt>
-<dd class="field-odd"><ul class="simple">
-<li><p><strong>X_train</strong> (<em>np.ndarray</em>) – Training data</p></li>
-<li><p><strong>y_train</strong> (<em>np.array</em>) – Mapped Training targets</p></li>
-<li><p><strong>eval_set</strong> (<em>list of tuple</em>) – List of eval tuple set (X, y)</p></li>
-<li><p><strong>weights</strong> (<em>either 0</em><em>, </em><em>1</em><em>, </em><em>dict</em><em> or </em><em>iterable</em>) – <p>if 0 (default) : no weights will be applied
-if 1 : classification only, will balanced class with inverse frequency
-if dict : keys are corresponding class values are sample weights
-if iterable : list or np array must be of length equal to nb elements</p>
-<blockquote>
-<div><p>in the training set</p>
-</div></blockquote>
-</p></li>
-<li><p><strong>batch_size</strong> (<em>int</em>) – how many samples per batch to load</p></li>
-<li><p><strong>num_workers</strong> (<em>int</em>) – how many subprocesses to use for data loading. 0 means that the data
-will be loaded in the main process</p></li>
-<li><p><strong>drop_last</strong> (<em>bool</em>) – set to True to drop the last incomplete batch, if the dataset size is not
-divisible by the batch size. If False and the size of dataset is not
-divisible by the batch size, then the last batch will be smaller</p></li>
-<li><p><strong>pin_memory</strong> (<em>bool</em>) – Whether to pin GPU memory during training</p></li>
-</ul>
-</dd>
-<dt class="field-even">Returns</dt>
-<dd class="field-even"><p><strong>train_dataloader, valid_dataloader</strong> – Training and validation dataloaders</p>
-</dd>
-<dt class="field-odd">Return type</dt>
-<dd class="field-odd"><p>torch.DataLoader, torch.DataLoader</p>
-</dd>
-</dl>
-</dd></dl>
-
-<dl class="function">
-<dt id="pytorch_tabnet.utils.create_explain_matrix">
-<code class="sig-prename descclassname">pytorch_tabnet.utils.</code><code class="sig-name descname">create_explain_matrix</code><span class="sig-paren">(</span><em class="sig-param">input_dim</em>, <em class="sig-param">cat_emb_dim</em>, <em class="sig-param">cat_idxs</em>, <em class="sig-param">post_embed_dim</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/utils.html#create_explain_matrix"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.utils.create_explain_matrix" title="Permalink to this definition">¶</a></dt>
-<dd><p>This is a computational trick.
-In order to rapidly sum importances from same embeddings
-to the initial index.</p>
-<dl class="field-list simple">
-<dt class="field-odd">Parameters</dt>
-<dd class="field-odd"><ul class="simple">
-<li><p><strong>input_dim</strong> (<em>int</em>) – Initial input dim</p></li>
-<li><p><strong>cat_emb_dim</strong> (<em>int</em><em> or </em><em>list of int</em>) – if int : size of embedding for all categorical feature
-if list of int : size of embedding for each categorical feature</p></li>
-<li><p><strong>cat_idxs</strong> (<em>list of int</em>) – Initial position of categorical features</p></li>
-<li><p><strong>post_embed_dim</strong> (<em>int</em>) – Post embedding inputs dimension</p></li>
-</ul>
-</dd>
-<dt class="field-even">Returns</dt>
-<dd class="field-even"><p><strong>reducing_matrix</strong> – Matrix of dim (post_embed_dim, input_dim)  to performe reduce</p>
-</dd>
-<dt class="field-odd">Return type</dt>
-<dd class="field-odd"><p>np.array</p>
-</dd>
-</dl>
-</dd></dl>
-
-<dl class="function">
-<dt id="pytorch_tabnet.utils.create_sampler">
-<code class="sig-prename descclassname">pytorch_tabnet.utils.</code><code class="sig-name descname">create_sampler</code><span class="sig-paren">(</span><em class="sig-param">weights</em>, <em class="sig-param">y_train</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/utils.html#create_sampler"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.utils.create_sampler" title="Permalink to this definition">¶</a></dt>
-<dd><p>This creates a sampler from the given weights</p>
-<dl class="field-list simple">
-<dt class="field-odd">Parameters</dt>
-<dd class="field-odd"><ul class="simple">
-<li><p><strong>weights</strong> (<em>either 0</em><em>, </em><em>1</em><em>, </em><em>dict</em><em> or </em><em>iterable</em>) – <p>if 0 (default) : no weights will be applied
-if 1 : classification only, will balanced class with inverse frequency
-if dict : keys are corresponding class values are sample weights
-if iterable : list or np array must be of length equal to nb elements</p>
-<blockquote>
-<div><p>in the training set</p>
-</div></blockquote>
-</p></li>
-<li><p><strong>y_train</strong> (<em>np.array</em>) – Training targets</p></li>
-</ul>
-</dd>
-</dl>
-</dd></dl>
-
-<dl class="function">
-<dt id="pytorch_tabnet.utils.define_device">
-<code class="sig-prename descclassname">pytorch_tabnet.utils.</code><code class="sig-name descname">define_device</code><span class="sig-paren">(</span><em class="sig-param">device_name</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/utils.html#define_device"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.utils.define_device" title="Permalink to this definition">¶</a></dt>
-<dd><p>Define the device to use during training and inference.
-If auto it will detect automatically whether to use cuda or cpu</p>
-<dl class="field-list simple">
-<dt class="field-odd">Parameters</dt>
-<dd class="field-odd"><p><strong>device_name</strong> (<em>str</em>) – Either “auto”, “cpu” or “cuda”</p>
-</dd>
-<dt class="field-even">Returns</dt>
-<dd class="field-even"><p>Either “cpu” or “cuda”</p>
-</dd>
-<dt class="field-odd">Return type</dt>
-<dd class="field-odd"><p>str</p>
-</dd>
-</dl>
-</dd></dl>
-
-<dl class="function">
-<dt id="pytorch_tabnet.utils.filter_weights">
-<code class="sig-prename descclassname">pytorch_tabnet.utils.</code><code class="sig-name descname">filter_weights</code><span class="sig-paren">(</span><em class="sig-param">weights</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/utils.html#filter_weights"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.utils.filter_weights" title="Permalink to this definition">¶</a></dt>
-<dd><p>This function makes sure that weights are in correct format for
-regression and multitask TabNet</p>
-<dl class="field-list simple">
-<dt class="field-odd">Parameters</dt>
-<dd class="field-odd"><p><strong>weights</strong> (<em>int</em><em>, </em><em>dict</em><em> or </em><em>list</em>) – Initial weights parameters given by user</p>
-</dd>
-<dt class="field-even">Returns</dt>
-<dd class="field-even"><p><strong>None</strong></p>
-</dd>
-<dt class="field-odd">Return type</dt>
-<dd class="field-odd"><p>This function will only throw an error if format is wrong</p>
-</dd>
-</dl>
-</dd></dl>
-
-<dl class="function">
-<dt id="pytorch_tabnet.utils.validate_eval_set">
-<code class="sig-prename descclassname">pytorch_tabnet.utils.</code><code class="sig-name descname">validate_eval_set</code><span class="sig-paren">(</span><em class="sig-param">eval_set</em>, <em class="sig-param">eval_name</em>, <em class="sig-param">X_train</em>, <em class="sig-param">y_train</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/utils.html#validate_eval_set"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.utils.validate_eval_set" title="Permalink to this definition">¶</a></dt>
-<dd><p>Check if the shapes of eval_set are compatible with (X_train, y_train).</p>
-<dl class="field-list simple">
-<dt class="field-odd">Parameters</dt>
-<dd class="field-odd"><ul class="simple">
-<li><p><strong>eval_set</strong> (<em>list of tuple</em>) – List of eval tuple set (X, y).
-The last one is used for early stopping</p></li>
-<li><p><strong>eval_name</strong> (<em>list of str</em>) – List of eval set names.</p></li>
-<li><p><strong>X_train</strong> (<em>np.ndarray</em>) – Train owned products</p></li>
-<li><p><strong>y_train</strong> (<em>np.array</em>) – Train targeted products</p></li>
-</ul>
-</dd>
-<dt class="field-even">Returns</dt>
-<dd class="field-even"><p><ul class="simple">
-<li><p><strong>eval_names</strong> (<em>list of str</em>) – Validated list of eval_names.</p></li>
-<li><p><strong>eval_set</strong> (<em>list of tuple</em>) – Validated list of eval_set.</p></li>
-</ul>
-</p>
-</dd>
-</dl>
-</dd></dl>
-
-</div>
-<div class="section" id="module-pytorch_tabnet.pretraining_utils">
+<section id="module-pytorch_tabnet.pretraining_utils">
 <span id="pytorch-tabnet-pretraining-utils-module"></span><h2>pytorch_tabnet.pretraining_utils module<a class="headerlink" href="#module-pytorch_tabnet.pretraining_utils" title="Permalink to this headline">¶</a></h2>
 <dl class="function">
 <dt id="pytorch_tabnet.pretraining_utils.create_dataloaders">
@@ -433,8 +239,30 @@ <h1>pytorch_tabnet package<a class="headerlink" href="#pytorch-tabnet-package" t
 </dl>
 </dd></dl>
 
-</div>
-<div class="section" id="module-pytorch_tabnet.tab_network">
+</section>
+<section id="module-pytorch_tabnet.augmentations">
+<span id="pytorch-tabnet-augmentations-module"></span><h2>pytorch_tabnet.augmentations module<a class="headerlink" href="#module-pytorch_tabnet.augmentations" title="Permalink to this headline">¶</a></h2>
+<dl class="class">
+<dt id="pytorch_tabnet.augmentations.ClassificationSMOTE">
+<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.augmentations.</code><code class="sig-name descname">ClassificationSMOTE</code><span class="sig-paren">(</span><em class="sig-param">device_name='auto'</em>, <em class="sig-param">p=0.8</em>, <em class="sig-param">alpha=0.5</em>, <em class="sig-param">beta=0.5</em>, <em class="sig-param">seed=0</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/augmentations.html#ClassificationSMOTE"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.augmentations.ClassificationSMOTE" title="Permalink to this definition">¶</a></dt>
+<dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">object</span></code></p>
+<p>Apply SMOTE for classification tasks.</p>
+<p>This will average a percentage p of the elements in the batch with other elements.
+The target will stay unchanged and keep the value of the most important row in the mix.</p>
+</dd></dl>
+
+<dl class="class">
+<dt id="pytorch_tabnet.augmentations.RegressionSMOTE">
+<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.augmentations.</code><code class="sig-name descname">RegressionSMOTE</code><span class="sig-paren">(</span><em class="sig-param">device_name='auto'</em>, <em class="sig-param">p=0.8</em>, <em class="sig-param">alpha=0.5</em>, <em class="sig-param">beta=0.5</em>, <em class="sig-param">seed=0</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/augmentations.html#RegressionSMOTE"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.augmentations.RegressionSMOTE" title="Permalink to this definition">¶</a></dt>
+<dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">object</span></code></p>
+<p>Apply SMOTE</p>
+<p>This will average a percentage p of the elements in the batch with other elements.
+The target will be averaged as well (this might work with binary classification
+and certain loss), following a beta distribution.</p>
+</dd></dl>
+
+</section>
+<section id="module-pytorch_tabnet.tab_network">
 <span id="pytorch-tabnet-tab-network-module"></span><h2>pytorch_tabnet.tab_network module<a class="headerlink" href="#module-pytorch_tabnet.tab_network" title="Permalink to this headline">¶</a></h2>
 <dl class="class">
 <dt id="pytorch_tabnet.tab_network.AttentiveTransformer">
@@ -602,7 +430,7 @@ <h1>pytorch_tabnet package<a class="headerlink" href="#pytorch-tabnet-package" t
 
 <dl class="class">
 <dt id="pytorch_tabnet.tab_network.TabNetDecoder">
-<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.tab_network.</code><code class="sig-name descname">TabNetDecoder</code><span class="sig-paren">(</span><em class="sig-param">input_dim</em>, <em class="sig-param">n_d=8</em>, <em class="sig-param">n_steps=3</em>, <em class="sig-param">n_independent=2</em>, <em class="sig-param">n_shared=2</em>, <em class="sig-param">virtual_batch_size=128</em>, <em class="sig-param">momentum=0.02</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/tab_network.html#TabNetDecoder"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.tab_network.TabNetDecoder" title="Permalink to this definition">¶</a></dt>
+<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.tab_network.</code><code class="sig-name descname">TabNetDecoder</code><span class="sig-paren">(</span><em class="sig-param">input_dim</em>, <em class="sig-param">n_d=8</em>, <em class="sig-param">n_steps=3</em>, <em class="sig-param">n_independent=1</em>, <em class="sig-param">n_shared=1</em>, <em class="sig-param">virtual_batch_size=128</em>, <em class="sig-param">momentum=0.02</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/tab_network.html#TabNetDecoder"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.tab_network.TabNetDecoder" title="Permalink to this definition">¶</a></dt>
 <dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">torch.nn.modules.module.Module</span></code></p>
 <dl class="method">
 <dt id="pytorch_tabnet.tab_network.TabNetDecoder.forward">
@@ -672,7 +500,7 @@ <h1>pytorch_tabnet package<a class="headerlink" href="#pytorch-tabnet-package" t
 
 <dl class="class">
 <dt id="pytorch_tabnet.tab_network.TabNetPretraining">
-<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.tab_network.</code><code class="sig-name descname">TabNetPretraining</code><span class="sig-paren">(</span><em class="sig-param">input_dim</em>, <em class="sig-param">pretraining_ratio=0.2</em>, <em class="sig-param">n_d=8</em>, <em class="sig-param">n_a=8</em>, <em class="sig-param">n_steps=3</em>, <em class="sig-param">gamma=1.3</em>, <em class="sig-param">cat_idxs=[]</em>, <em class="sig-param">cat_dims=[]</em>, <em class="sig-param">cat_emb_dim=1</em>, <em class="sig-param">n_independent=2</em>, <em class="sig-param">n_shared=2</em>, <em class="sig-param">epsilon=1e-15</em>, <em class="sig-param">virtual_batch_size=128</em>, <em class="sig-param">momentum=0.02</em>, <em class="sig-param">mask_type='sparsemax'</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/tab_network.html#TabNetPretraining"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.tab_network.TabNetPretraining" title="Permalink to this definition">¶</a></dt>
+<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.tab_network.</code><code class="sig-name descname">TabNetPretraining</code><span class="sig-paren">(</span><em class="sig-param">input_dim</em>, <em class="sig-param">pretraining_ratio=0.2</em>, <em class="sig-param">n_d=8</em>, <em class="sig-param">n_a=8</em>, <em class="sig-param">n_steps=3</em>, <em class="sig-param">gamma=1.3</em>, <em class="sig-param">cat_idxs=[]</em>, <em class="sig-param">cat_dims=[]</em>, <em class="sig-param">cat_emb_dim=1</em>, <em class="sig-param">n_independent=2</em>, <em class="sig-param">n_shared=2</em>, <em class="sig-param">epsilon=1e-15</em>, <em class="sig-param">virtual_batch_size=128</em>, <em class="sig-param">momentum=0.02</em>, <em class="sig-param">mask_type='sparsemax'</em>, <em class="sig-param">n_shared_decoder=1</em>, <em class="sig-param">n_indep_decoder=1</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/tab_network.html#TabNetPretraining"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.tab_network.TabNetPretraining" title="Permalink to this definition">¶</a></dt>
 <dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">torch.nn.modules.module.Module</span></code></p>
 <dl class="method">
 <dt id="pytorch_tabnet.tab_network.TabNetPretraining.forward">
@@ -702,68 +530,259 @@ <h1>pytorch_tabnet package<a class="headerlink" href="#pytorch-tabnet-package" t
 <code class="sig-prename descclassname">pytorch_tabnet.tab_network.</code><code class="sig-name descname">initialize_non_glu</code><span class="sig-paren">(</span><em class="sig-param">module</em>, <em class="sig-param">input_dim</em>, <em class="sig-param">output_dim</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/tab_network.html#initialize_non_glu"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.tab_network.initialize_non_glu" title="Permalink to this definition">¶</a></dt>
 <dd></dd></dl>
 
-</div>
-<div class="section" id="module-pytorch_tabnet.multitask">
-<span id="pytorch-tabnet-multitask-module"></span><h2>pytorch_tabnet.multitask module<a class="headerlink" href="#module-pytorch_tabnet.multitask" title="Permalink to this headline">¶</a></h2>
+</section>
+<section id="module-pytorch_tabnet.metrics">
+<span id="pytorch-tabnet-metrics-module"></span><h2>pytorch_tabnet.metrics module<a class="headerlink" href="#module-pytorch_tabnet.metrics" title="Permalink to this headline">¶</a></h2>
 <dl class="class">
-<dt id="pytorch_tabnet.multitask.TabNetMultiTaskClassifier">
-<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.multitask.</code><code class="sig-name descname">TabNetMultiTaskClassifier</code><span class="sig-paren">(</span><em class="sig-param">n_d: int = 8</em>, <em class="sig-param">n_a: int = 8</em>, <em class="sig-param">n_steps: int = 3</em>, <em class="sig-param">gamma: float = 1.3</em>, <em class="sig-param">cat_idxs: List[int] = &lt;factory&gt;</em>, <em class="sig-param">cat_dims: List[int] = &lt;factory&gt;</em>, <em class="sig-param">cat_emb_dim: int = 1</em>, <em class="sig-param">n_independent: int = 2</em>, <em class="sig-param">n_shared: int = 2</em>, <em class="sig-param">epsilon: float = 1e-15</em>, <em class="sig-param">momentum: float = 0.02</em>, <em class="sig-param">lambda_sparse: float = 0.001</em>, <em class="sig-param">seed: int = 0</em>, <em class="sig-param">clip_value: int = 1</em>, <em class="sig-param">verbose: int = 1</em>, <em class="sig-param">optimizer_fn: Any = &lt;class 'torch.optim.adam.Adam'&gt;</em>, <em class="sig-param">optimizer_params: Dict = &lt;factory&gt;</em>, <em class="sig-param">scheduler_fn: Any = None</em>, <em class="sig-param">scheduler_params: Dict = &lt;factory&gt;</em>, <em class="sig-param">mask_type: str = 'sparsemax'</em>, <em class="sig-param">input_dim: int = None</em>, <em class="sig-param">output_dim: int = None</em>, <em class="sig-param">device_name: str = 'auto'</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/multitask.html#TabNetMultiTaskClassifier"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.multitask.TabNetMultiTaskClassifier" title="Permalink to this definition">¶</a></dt>
-<dd><p>Bases: <a class="reference internal" href="#pytorch_tabnet.abstract_model.TabModel" title="pytorch_tabnet.abstract_model.TabModel"><code class="xref py py-class docutils literal notranslate"><span class="pre">pytorch_tabnet.abstract_model.TabModel</span></code></a></p>
-<dl class="attribute">
-<dt id="pytorch_tabnet.multitask.TabNetMultiTaskClassifier.cat_dims">
-<code class="sig-name descname">cat_dims</code><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.multitask.TabNetMultiTaskClassifier.cat_dims" title="Permalink to this definition">¶</a></dt>
-<dd></dd></dl>
+<dt id="pytorch_tabnet.metrics.AUC">
+<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.metrics.</code><code class="sig-name descname">AUC</code><a class="reference internal" href="../_modules/pytorch_tabnet/metrics.html#AUC"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.metrics.AUC" title="Permalink to this definition">¶</a></dt>
+<dd><p>Bases: <a class="reference internal" href="#pytorch_tabnet.metrics.Metric" title="pytorch_tabnet.metrics.Metric"><code class="xref py py-class docutils literal notranslate"><span class="pre">pytorch_tabnet.metrics.Metric</span></code></a></p>
+<p>AUC.</p>
+</dd></dl>
 
-<dl class="attribute">
-<dt id="pytorch_tabnet.multitask.TabNetMultiTaskClassifier.cat_idxs">
-<code class="sig-name descname">cat_idxs</code><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.multitask.TabNetMultiTaskClassifier.cat_idxs" title="Permalink to this definition">¶</a></dt>
-<dd></dd></dl>
+<dl class="class">
+<dt id="pytorch_tabnet.metrics.Accuracy">
+<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.metrics.</code><code class="sig-name descname">Accuracy</code><a class="reference internal" href="../_modules/pytorch_tabnet/metrics.html#Accuracy"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.metrics.Accuracy" title="Permalink to this definition">¶</a></dt>
+<dd><p>Bases: <a class="reference internal" href="#pytorch_tabnet.metrics.Metric" title="pytorch_tabnet.metrics.Metric"><code class="xref py py-class docutils literal notranslate"><span class="pre">pytorch_tabnet.metrics.Metric</span></code></a></p>
+<p>Accuracy.</p>
+</dd></dl>
 
-<dl class="method">
-<dt id="pytorch_tabnet.multitask.TabNetMultiTaskClassifier.compute_loss">
-<code class="sig-name descname">compute_loss</code><span class="sig-paren">(</span><em class="sig-param">y_pred</em>, <em class="sig-param">y_true</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/multitask.html#TabNetMultiTaskClassifier.compute_loss"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.multitask.TabNetMultiTaskClassifier.compute_loss" title="Permalink to this definition">¶</a></dt>
-<dd><p>Computes the loss according to network output and targets</p>
-<dl class="field-list simple">
-<dt class="field-odd">Parameters</dt>
-<dd class="field-odd"><ul class="simple">
-<li><p><strong>y_pred</strong> (<em>list of tensors</em>) – Output of network</p></li>
-<li><p><strong>y_true</strong> (<em>LongTensor</em>) – Targets label encoded</p></li>
+<dl class="class">
+<dt id="pytorch_tabnet.metrics.BalancedAccuracy">
+<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.metrics.</code><code class="sig-name descname">BalancedAccuracy</code><a class="reference internal" href="../_modules/pytorch_tabnet/metrics.html#BalancedAccuracy"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.metrics.BalancedAccuracy" title="Permalink to this definition">¶</a></dt>
+<dd><p>Bases: <a class="reference internal" href="#pytorch_tabnet.metrics.Metric" title="pytorch_tabnet.metrics.Metric"><code class="xref py py-class docutils literal notranslate"><span class="pre">pytorch_tabnet.metrics.Metric</span></code></a></p>
+<p>Balanced Accuracy.</p>
+</dd></dl>
+
+<dl class="class">
+<dt id="pytorch_tabnet.metrics.LogLoss">
+<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.metrics.</code><code class="sig-name descname">LogLoss</code><a class="reference internal" href="../_modules/pytorch_tabnet/metrics.html#LogLoss"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.metrics.LogLoss" title="Permalink to this definition">¶</a></dt>
+<dd><p>Bases: <a class="reference internal" href="#pytorch_tabnet.metrics.Metric" title="pytorch_tabnet.metrics.Metric"><code class="xref py py-class docutils literal notranslate"><span class="pre">pytorch_tabnet.metrics.Metric</span></code></a></p>
+<p>LogLoss.</p>
+</dd></dl>
+
+<dl class="class">
+<dt id="pytorch_tabnet.metrics.MAE">
+<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.metrics.</code><code class="sig-name descname">MAE</code><a class="reference internal" href="../_modules/pytorch_tabnet/metrics.html#MAE"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.metrics.MAE" title="Permalink to this definition">¶</a></dt>
+<dd><p>Bases: <a class="reference internal" href="#pytorch_tabnet.metrics.Metric" title="pytorch_tabnet.metrics.Metric"><code class="xref py py-class docutils literal notranslate"><span class="pre">pytorch_tabnet.metrics.Metric</span></code></a></p>
+<p>Mean Absolute Error.</p>
+</dd></dl>
+
+<dl class="class">
+<dt id="pytorch_tabnet.metrics.MSE">
+<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.metrics.</code><code class="sig-name descname">MSE</code><a class="reference internal" href="../_modules/pytorch_tabnet/metrics.html#MSE"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.metrics.MSE" title="Permalink to this definition">¶</a></dt>
+<dd><p>Bases: <a class="reference internal" href="#pytorch_tabnet.metrics.Metric" title="pytorch_tabnet.metrics.Metric"><code class="xref py py-class docutils literal notranslate"><span class="pre">pytorch_tabnet.metrics.Metric</span></code></a></p>
+<p>Mean Squared Error.</p>
+</dd></dl>
+
+<dl class="class">
+<dt id="pytorch_tabnet.metrics.Metric">
+<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.metrics.</code><code class="sig-name descname">Metric</code><a class="reference internal" href="../_modules/pytorch_tabnet/metrics.html#Metric"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.metrics.Metric" title="Permalink to this definition">¶</a></dt>
+<dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">object</span></code></p>
+<dl class="method">
+<dt id="pytorch_tabnet.metrics.Metric.get_metrics_by_names">
+<em class="property">classmethod </em><code class="sig-name descname">get_metrics_by_names</code><span class="sig-paren">(</span><em class="sig-param">names</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/metrics.html#Metric.get_metrics_by_names"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.metrics.Metric.get_metrics_by_names" title="Permalink to this definition">¶</a></dt>
+<dd><p>Get list of metric classes.</p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters</dt>
+<dd class="field-odd"><ul class="simple">
+<li><p><strong>cls</strong> (<a class="reference internal" href="#pytorch_tabnet.metrics.Metric" title="pytorch_tabnet.metrics.Metric"><em>Metric</em></a>) – Metric class.</p></li>
+<li><p><strong>names</strong> (<em>list</em>) – List of metric names.</p></li>
 </ul>
 </dd>
 <dt class="field-even">Returns</dt>
-<dd class="field-even"><p><strong>loss</strong> – output of loss function(s)</p>
+<dd class="field-even"><p><strong>metrics</strong> – List of metric classes.</p>
 </dd>
 <dt class="field-odd">Return type</dt>
-<dd class="field-odd"><p>torch.Tensor</p>
+<dd class="field-odd"><p>list</p>
 </dd>
 </dl>
 </dd></dl>
 
+</dd></dl>
+
+<dl class="class">
+<dt id="pytorch_tabnet.metrics.MetricContainer">
+<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.metrics.</code><code class="sig-name descname">MetricContainer</code><span class="sig-paren">(</span><em class="sig-param">metric_names: List[str], prefix: str = ''</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/metrics.html#MetricContainer"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.metrics.MetricContainer" title="Permalink to this definition">¶</a></dt>
+<dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">object</span></code></p>
+<p>Container holding a list of metrics.</p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters</dt>
+<dd class="field-odd"><ul class="simple">
+<li><p><strong>metric_names</strong> (<em>list of str</em>) – List of metric names.</p></li>
+<li><p><strong>prefix</strong> (<em>str</em>) – Prefix of metric names.</p></li>
+</ul>
+</dd>
+</dl>
 <dl class="attribute">
-<dt id="pytorch_tabnet.multitask.TabNetMultiTaskClassifier.optimizer_params">
-<code class="sig-name descname">optimizer_params</code><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.multitask.TabNetMultiTaskClassifier.optimizer_params" title="Permalink to this definition">¶</a></dt>
+<dt id="pytorch_tabnet.metrics.MetricContainer.metric_names">
+<code class="sig-name descname">metric_names</code><em class="property">: List[str]</em><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.metrics.MetricContainer.metric_names" title="Permalink to this definition">¶</a></dt>
+<dd></dd></dl>
+
+<dl class="attribute">
+<dt id="pytorch_tabnet.metrics.MetricContainer.prefix">
+<code class="sig-name descname">prefix</code><em class="property">: str</em><em class="property"> = ''</em><a class="headerlink" href="#pytorch_tabnet.metrics.MetricContainer.prefix" title="Permalink to this definition">¶</a></dt>
+<dd></dd></dl>
+
+</dd></dl>
+
+<dl class="class">
+<dt id="pytorch_tabnet.metrics.RMSE">
+<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.metrics.</code><code class="sig-name descname">RMSE</code><a class="reference internal" href="../_modules/pytorch_tabnet/metrics.html#RMSE"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.metrics.RMSE" title="Permalink to this definition">¶</a></dt>
+<dd><p>Bases: <a class="reference internal" href="#pytorch_tabnet.metrics.Metric" title="pytorch_tabnet.metrics.Metric"><code class="xref py py-class docutils literal notranslate"><span class="pre">pytorch_tabnet.metrics.Metric</span></code></a></p>
+<p>Root Mean Squared Error.</p>
+</dd></dl>
+
+<dl class="class">
+<dt id="pytorch_tabnet.metrics.RMSLE">
+<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.metrics.</code><code class="sig-name descname">RMSLE</code><a class="reference internal" href="../_modules/pytorch_tabnet/metrics.html#RMSLE"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.metrics.RMSLE" title="Permalink to this definition">¶</a></dt>
+<dd><p>Bases: <a class="reference internal" href="#pytorch_tabnet.metrics.Metric" title="pytorch_tabnet.metrics.Metric"><code class="xref py py-class docutils literal notranslate"><span class="pre">pytorch_tabnet.metrics.Metric</span></code></a></p>
+<p>Mean squared logarithmic error regression loss.
+Scikit-implementation:
+<a class="reference external" href="https://scikit-learn.org/stable/modules/generated/sklearn.metrics.mean_squared_log_error.html">https://scikit-learn.org/stable/modules/generated/sklearn.metrics.mean_squared_log_error.html</a>
+Note: In order to avoid error, negative predictions are clipped to 0.
+This means that you should clip negative predictions manually after calling predict.</p>
+</dd></dl>
+
+<dl class="class">
+<dt id="pytorch_tabnet.metrics.UnsupMetricContainer">
+<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.metrics.</code><code class="sig-name descname">UnsupMetricContainer</code><span class="sig-paren">(</span><em class="sig-param">metric_names: List[str], prefix: str = ''</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/metrics.html#UnsupMetricContainer"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.metrics.UnsupMetricContainer" title="Permalink to this definition">¶</a></dt>
+<dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">object</span></code></p>
+<p>Container holding a list of metrics.</p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters</dt>
+<dd class="field-odd"><ul class="simple">
+<li><p><strong>y_pred</strong> (<em>torch.Tensor</em><em> or </em><em>np.array</em>) – Reconstructed prediction (with embeddings)</p></li>
+<li><p><strong>embedded_x</strong> (<em>torch.Tensor</em>) – Original input embedded by network</p></li>
+<li><p><strong>obf_vars</strong> (<em>torch.Tensor</em>) – Binary mask for obfuscated variables.
+1 means the variables was obfuscated so reconstruction is based on this.</p></li>
+</ul>
+</dd>
+</dl>
+<dl class="attribute">
+<dt id="pytorch_tabnet.metrics.UnsupMetricContainer.metric_names">
+<code class="sig-name descname">metric_names</code><em class="property">: List[str]</em><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.metrics.UnsupMetricContainer.metric_names" title="Permalink to this definition">¶</a></dt>
+<dd></dd></dl>
+
+<dl class="attribute">
+<dt id="pytorch_tabnet.metrics.UnsupMetricContainer.prefix">
+<code class="sig-name descname">prefix</code><em class="property">: str</em><em class="property"> = ''</em><a class="headerlink" href="#pytorch_tabnet.metrics.UnsupMetricContainer.prefix" title="Permalink to this definition">¶</a></dt>
+<dd></dd></dl>
+
+</dd></dl>
+
+<dl class="function">
+<dt id="pytorch_tabnet.metrics.UnsupervisedLoss">
+<code class="sig-prename descclassname">pytorch_tabnet.metrics.</code><code class="sig-name descname">UnsupervisedLoss</code><span class="sig-paren">(</span><em class="sig-param">y_pred</em>, <em class="sig-param">embedded_x</em>, <em class="sig-param">obf_vars</em>, <em class="sig-param">eps=1e-09</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/metrics.html#UnsupervisedLoss"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.metrics.UnsupervisedLoss" title="Permalink to this definition">¶</a></dt>
+<dd><p>Implements unsupervised loss function.
+This differs from orginal paper as it’s scaled to be batch size independent
+and number of features reconstructed independent (by taking the mean)</p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters</dt>
+<dd class="field-odd"><ul class="simple">
+<li><p><strong>y_pred</strong> (<em>torch.Tensor</em><em> or </em><em>np.array</em>) – Reconstructed prediction (with embeddings)</p></li>
+<li><p><strong>embedded_x</strong> (<em>torch.Tensor</em>) – Original input embedded by network</p></li>
+<li><p><strong>obf_vars</strong> (<em>torch.Tensor</em>) – Binary mask for obfuscated variables.
+1 means the variable was obfuscated so reconstruction is based on this.</p></li>
+<li><p><strong>eps</strong> (<em>float</em>) – A small floating point to avoid ZeroDivisionError
+This can happen in degenerated case when a feature has only one value</p></li>
+</ul>
+</dd>
+<dt class="field-even">Returns</dt>
+<dd class="field-even"><p><strong>loss</strong> – Unsupervised loss, average value over batch samples.</p>
+</dd>
+<dt class="field-odd">Return type</dt>
+<dd class="field-odd"><p>torch float</p>
+</dd>
+</dl>
+</dd></dl>
+
+<dl class="function">
+<dt id="pytorch_tabnet.metrics.UnsupervisedLossNumpy">
+<code class="sig-prename descclassname">pytorch_tabnet.metrics.</code><code class="sig-name descname">UnsupervisedLossNumpy</code><span class="sig-paren">(</span><em class="sig-param">y_pred</em>, <em class="sig-param">embedded_x</em>, <em class="sig-param">obf_vars</em>, <em class="sig-param">eps=1e-09</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/metrics.html#UnsupervisedLossNumpy"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.metrics.UnsupervisedLossNumpy" title="Permalink to this definition">¶</a></dt>
+<dd></dd></dl>
+
+<dl class="class">
+<dt id="pytorch_tabnet.metrics.UnsupervisedMetric">
+<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.metrics.</code><code class="sig-name descname">UnsupervisedMetric</code><a class="reference internal" href="../_modules/pytorch_tabnet/metrics.html#UnsupervisedMetric"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.metrics.UnsupervisedMetric" title="Permalink to this definition">¶</a></dt>
+<dd><p>Bases: <a class="reference internal" href="#pytorch_tabnet.metrics.Metric" title="pytorch_tabnet.metrics.Metric"><code class="xref py py-class docutils literal notranslate"><span class="pre">pytorch_tabnet.metrics.Metric</span></code></a></p>
+<p>Unsupervised metric</p>
+</dd></dl>
+
+<dl class="class">
+<dt id="pytorch_tabnet.metrics.UnsupervisedNumpyMetric">
+<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.metrics.</code><code class="sig-name descname">UnsupervisedNumpyMetric</code><a class="reference internal" href="../_modules/pytorch_tabnet/metrics.html#UnsupervisedNumpyMetric"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.metrics.UnsupervisedNumpyMetric" title="Permalink to this definition">¶</a></dt>
+<dd><p>Bases: <a class="reference internal" href="#pytorch_tabnet.metrics.Metric" title="pytorch_tabnet.metrics.Metric"><code class="xref py py-class docutils literal notranslate"><span class="pre">pytorch_tabnet.metrics.Metric</span></code></a></p>
+<p>Unsupervised metric</p>
+</dd></dl>
+
+<dl class="function">
+<dt id="pytorch_tabnet.metrics.check_metrics">
+<code class="sig-prename descclassname">pytorch_tabnet.metrics.</code><code class="sig-name descname">check_metrics</code><span class="sig-paren">(</span><em class="sig-param">metrics</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/metrics.html#check_metrics"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.metrics.check_metrics" title="Permalink to this definition">¶</a></dt>
+<dd><p>Check if custom metrics are provided.</p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters</dt>
+<dd class="field-odd"><p><strong>metrics</strong> (<em>list of str</em><em> or </em><em>classes</em>) – List with built-in metrics (str) or custom metrics (classes).</p>
+</dd>
+<dt class="field-even">Returns</dt>
+<dd class="field-even"><p><strong>val_metrics</strong> – List of metric names.</p>
+</dd>
+<dt class="field-odd">Return type</dt>
+<dd class="field-odd"><p>list of str</p>
+</dd>
+</dl>
+</dd></dl>
+
+</section>
+<section id="module-pytorch_tabnet.tab_model">
+<span id="pytorch-tabnet-tab-model-module"></span><h2>pytorch_tabnet.tab_model module<a class="headerlink" href="#module-pytorch_tabnet.tab_model" title="Permalink to this headline">¶</a></h2>
+<dl class="class">
+<dt id="pytorch_tabnet.tab_model.TabNetClassifier">
+<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.tab_model.</code><code class="sig-name descname">TabNetClassifier</code><span class="sig-paren">(</span><em class="sig-param">n_d: int = 8</em>, <em class="sig-param">n_a: int = 8</em>, <em class="sig-param">n_steps: int = 3</em>, <em class="sig-param">gamma: float = 1.3</em>, <em class="sig-param">cat_idxs: List[int] = &lt;factory&gt;</em>, <em class="sig-param">cat_dims: List[int] = &lt;factory&gt;</em>, <em class="sig-param">cat_emb_dim: int = 1</em>, <em class="sig-param">n_independent: int = 2</em>, <em class="sig-param">n_shared: int = 2</em>, <em class="sig-param">epsilon: float = 1e-15</em>, <em class="sig-param">momentum: float = 0.02</em>, <em class="sig-param">lambda_sparse: float = 0.001</em>, <em class="sig-param">seed: int = 0</em>, <em class="sig-param">clip_value: int = 1</em>, <em class="sig-param">verbose: int = 1</em>, <em class="sig-param">optimizer_fn: Any = &lt;class 'torch.optim.adam.Adam'&gt;</em>, <em class="sig-param">optimizer_params: Dict = &lt;factory&gt;</em>, <em class="sig-param">scheduler_fn: Any = None</em>, <em class="sig-param">scheduler_params: Dict = &lt;factory&gt;</em>, <em class="sig-param">mask_type: str = 'sparsemax'</em>, <em class="sig-param">input_dim: int = None</em>, <em class="sig-param">output_dim: int = None</em>, <em class="sig-param">device_name: str = 'auto'</em>, <em class="sig-param">n_shared_decoder: int = 1</em>, <em class="sig-param">n_indep_decoder: int = 1</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/tab_model.html#TabNetClassifier"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetClassifier" title="Permalink to this definition">¶</a></dt>
+<dd><p>Bases: <a class="reference internal" href="#pytorch_tabnet.abstract_model.TabModel" title="pytorch_tabnet.abstract_model.TabModel"><code class="xref py py-class docutils literal notranslate"><span class="pre">pytorch_tabnet.abstract_model.TabModel</span></code></a></p>
+<dl class="attribute">
+<dt id="pytorch_tabnet.tab_model.TabNetClassifier.cat_dims">
+<code class="sig-name descname">cat_dims</code><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetClassifier.cat_dims" title="Permalink to this definition">¶</a></dt>
+<dd></dd></dl>
+
+<dl class="attribute">
+<dt id="pytorch_tabnet.tab_model.TabNetClassifier.cat_idxs">
+<code class="sig-name descname">cat_idxs</code><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetClassifier.cat_idxs" title="Permalink to this definition">¶</a></dt>
 <dd></dd></dl>
 
 <dl class="method">
-<dt id="pytorch_tabnet.multitask.TabNetMultiTaskClassifier.predict">
-<code class="sig-name descname">predict</code><span class="sig-paren">(</span><em class="sig-param">X</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/multitask.html#TabNetMultiTaskClassifier.predict"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.multitask.TabNetMultiTaskClassifier.predict" title="Permalink to this definition">¶</a></dt>
-<dd><p>Make predictions on a batch (valid)</p>
+<dt id="pytorch_tabnet.tab_model.TabNetClassifier.compute_loss">
+<code class="sig-name descname">compute_loss</code><span class="sig-paren">(</span><em class="sig-param">y_pred</em>, <em class="sig-param">y_true</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/tab_model.html#TabNetClassifier.compute_loss"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetClassifier.compute_loss" title="Permalink to this definition">¶</a></dt>
+<dd><p>Compute the loss.</p>
 <dl class="field-list simple">
 <dt class="field-odd">Parameters</dt>
-<dd class="field-odd"><p><strong>X</strong> (a :tensor: <cite>torch.Tensor</cite>) – Input data</p>
+<dd class="field-odd"><ul class="simple">
+<li><p><strong>y_score</strong> (a :tensor: <cite>torch.Tensor</cite>) – Score matrix</p></li>
+<li><p><strong>y_true</strong> (a :tensor: <cite>torch.Tensor</cite>) – Target matrix</p></li>
+</ul>
 </dd>
 <dt class="field-even">Returns</dt>
-<dd class="field-even"><p><strong>results</strong> – Predictions of the most probable class</p>
+<dd class="field-even"><p>Loss value</p>
 </dd>
 <dt class="field-odd">Return type</dt>
-<dd class="field-odd"><p>np.array</p>
+<dd class="field-odd"><p>float</p>
 </dd>
 </dl>
 </dd></dl>
 
+<dl class="attribute">
+<dt id="pytorch_tabnet.tab_model.TabNetClassifier.optimizer_params">
+<code class="sig-name descname">optimizer_params</code><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetClassifier.optimizer_params" title="Permalink to this definition">¶</a></dt>
+<dd></dd></dl>
+
 <dl class="method">
-<dt id="pytorch_tabnet.multitask.TabNetMultiTaskClassifier.predict_proba">
-<code class="sig-name descname">predict_proba</code><span class="sig-paren">(</span><em class="sig-param">X</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/multitask.html#TabNetMultiTaskClassifier.predict_proba"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.multitask.TabNetMultiTaskClassifier.predict_proba" title="Permalink to this definition">¶</a></dt>
+<dt id="pytorch_tabnet.tab_model.TabNetClassifier.predict_func">
+<code class="sig-name descname">predict_func</code><span class="sig-paren">(</span><em class="sig-param">outputs</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/tab_model.html#TabNetClassifier.predict_func"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetClassifier.predict_func" title="Permalink to this definition">¶</a></dt>
+<dd></dd></dl>
+
+<dl class="method">
+<dt id="pytorch_tabnet.tab_model.TabNetClassifier.predict_proba">
+<code class="sig-name descname">predict_proba</code><span class="sig-paren">(</span><em class="sig-param">X</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/tab_model.html#TabNetClassifier.predict_proba"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetClassifier.predict_proba" title="Permalink to this definition">¶</a></dt>
 <dd><p>Make predictions for classification on a batch (valid)</p>
 <dl class="field-list simple">
 <dt class="field-odd">Parameters</dt>
@@ -773,14 +792,14 @@ <h1>pytorch_tabnet package<a class="headerlink" href="#pytorch-tabnet-package" t
 <dd class="field-even"><p><strong>res</strong></p>
 </dd>
 <dt class="field-odd">Return type</dt>
-<dd class="field-odd"><p>list of np.ndarray</p>
+<dd class="field-odd"><p>np.ndarray</p>
 </dd>
 </dl>
 </dd></dl>
 
 <dl class="method">
-<dt id="pytorch_tabnet.multitask.TabNetMultiTaskClassifier.prepare_target">
-<code class="sig-name descname">prepare_target</code><span class="sig-paren">(</span><em class="sig-param">y</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/multitask.html#TabNetMultiTaskClassifier.prepare_target"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.multitask.TabNetMultiTaskClassifier.prepare_target" title="Permalink to this definition">¶</a></dt>
+<dt id="pytorch_tabnet.tab_model.TabNetClassifier.prepare_target">
+<code class="sig-name descname">prepare_target</code><span class="sig-paren">(</span><em class="sig-param">y</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/tab_model.html#TabNetClassifier.prepare_target"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetClassifier.prepare_target" title="Permalink to this definition">¶</a></dt>
 <dd><p>Prepare target before training.</p>
 <dl class="field-list simple">
 <dt class="field-odd">Parameters</dt>
@@ -796,18 +815,18 @@ <h1>pytorch_tabnet package<a class="headerlink" href="#pytorch-tabnet-package" t
 </dd></dl>
 
 <dl class="attribute">
-<dt id="pytorch_tabnet.multitask.TabNetMultiTaskClassifier.scheduler_params">
-<code class="sig-name descname">scheduler_params</code><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.multitask.TabNetMultiTaskClassifier.scheduler_params" title="Permalink to this definition">¶</a></dt>
+<dt id="pytorch_tabnet.tab_model.TabNetClassifier.scheduler_params">
+<code class="sig-name descname">scheduler_params</code><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetClassifier.scheduler_params" title="Permalink to this definition">¶</a></dt>
 <dd></dd></dl>
 
 <dl class="method">
-<dt id="pytorch_tabnet.multitask.TabNetMultiTaskClassifier.stack_batches">
-<code class="sig-name descname">stack_batches</code><span class="sig-paren">(</span><em class="sig-param">list_y_true</em>, <em class="sig-param">list_y_score</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/multitask.html#TabNetMultiTaskClassifier.stack_batches"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.multitask.TabNetMultiTaskClassifier.stack_batches" title="Permalink to this definition">¶</a></dt>
+<dt id="pytorch_tabnet.tab_model.TabNetClassifier.stack_batches">
+<code class="sig-name descname">stack_batches</code><span class="sig-paren">(</span><em class="sig-param">list_y_true</em>, <em class="sig-param">list_y_score</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/tab_model.html#TabNetClassifier.stack_batches"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetClassifier.stack_batches" title="Permalink to this definition">¶</a></dt>
 <dd></dd></dl>
 
 <dl class="method">
-<dt id="pytorch_tabnet.multitask.TabNetMultiTaskClassifier.update_fit_params">
-<code class="sig-name descname">update_fit_params</code><span class="sig-paren">(</span><em class="sig-param">X_train</em>, <em class="sig-param">y_train</em>, <em class="sig-param">eval_set</em>, <em class="sig-param">weights</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/multitask.html#TabNetMultiTaskClassifier.update_fit_params"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.multitask.TabNetMultiTaskClassifier.update_fit_params" title="Permalink to this definition">¶</a></dt>
+<dt id="pytorch_tabnet.tab_model.TabNetClassifier.update_fit_params">
+<code class="sig-name descname">update_fit_params</code><span class="sig-paren">(</span><em class="sig-param">X_train</em>, <em class="sig-param">y_train</em>, <em class="sig-param">eval_set</em>, <em class="sig-param">weights</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/tab_model.html#TabNetClassifier.update_fit_params"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetClassifier.update_fit_params" title="Permalink to this definition">¶</a></dt>
 <dd><p>Set attributes relative to fit function.</p>
 <dl class="field-list simple">
 <dt class="field-odd">Parameters</dt>
@@ -822,254 +841,42 @@ <h1>pytorch_tabnet package<a class="headerlink" href="#pytorch-tabnet-package" t
 </dl>
 </dd></dl>
 
-</dd></dl>
-
-</div>
-<div class="section" id="module-pytorch_tabnet.multiclass_utils">
-<span id="pytorch-tabnet-multiclass-utils-module"></span><h2>pytorch_tabnet.multiclass_utils module<a class="headerlink" href="#module-pytorch_tabnet.multiclass_utils" title="Permalink to this headline">¶</a></h2>
-<div class="section" id="multi-class-multi-label-utility-function">
-<h3>Multi-class / multi-label utility function<a class="headerlink" href="#multi-class-multi-label-utility-function" title="Permalink to this headline">¶</a></h3>
-<dl class="function">
-<dt id="pytorch_tabnet.multiclass_utils.assert_all_finite">
-<code class="sig-prename descclassname">pytorch_tabnet.multiclass_utils.</code><code class="sig-name descname">assert_all_finite</code><span class="sig-paren">(</span><em class="sig-param">X</em>, <em class="sig-param">allow_nan=False</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/multiclass_utils.html#assert_all_finite"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.multiclass_utils.assert_all_finite" title="Permalink to this definition">¶</a></dt>
-<dd><p>Throw a ValueError if X contains NaN or infinity.</p>
+<dl class="method">
+<dt id="pytorch_tabnet.tab_model.TabNetClassifier.weight_updater">
+<code class="sig-name descname">weight_updater</code><span class="sig-paren">(</span><em class="sig-param">weights</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/tab_model.html#TabNetClassifier.weight_updater"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetClassifier.weight_updater" title="Permalink to this definition">¶</a></dt>
+<dd><p>Updates weights dictionary according to target_mapper.</p>
 <dl class="field-list simple">
 <dt class="field-odd">Parameters</dt>
-<dd class="field-odd"><ul class="simple">
-<li><p><strong>X</strong> (<em>array</em><em> or </em><em>sparse matrix</em>) – </p></li>
-<li><p><strong>allow_nan</strong> (<em>bool</em>) – </p></li>
-</ul>
-</dd>
-</dl>
-</dd></dl>
-
-<dl class="function">
-<dt id="pytorch_tabnet.multiclass_utils.check_classification_targets">
-<code class="sig-prename descclassname">pytorch_tabnet.multiclass_utils.</code><code class="sig-name descname">check_classification_targets</code><span class="sig-paren">(</span><em class="sig-param">y</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/multiclass_utils.html#check_classification_targets"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.multiclass_utils.check_classification_targets" title="Permalink to this definition">¶</a></dt>
-<dd><p>Ensure that target y is of a non-regression type.</p>
-<dl class="simple">
-<dt>Only the following target types (as defined in type_of_target) are allowed:</dt><dd><p>‘binary’, ‘multiclass’, ‘multiclass-multioutput’,
-‘multilabel-indicator’, ‘multilabel-sequences’</p>
-</dd>
-</dl>
-<dl class="field-list simple">
-<dt class="field-odd">Parameters</dt>
-<dd class="field-odd"><p><strong>y</strong> (<em>array-like</em>) – </p>
-</dd>
-</dl>
-</dd></dl>
-
-<dl class="function">
-<dt id="pytorch_tabnet.multiclass_utils.check_output_dim">
-<code class="sig-prename descclassname">pytorch_tabnet.multiclass_utils.</code><code class="sig-name descname">check_output_dim</code><span class="sig-paren">(</span><em class="sig-param">labels</em>, <em class="sig-param">y</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/multiclass_utils.html#check_output_dim"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.multiclass_utils.check_output_dim" title="Permalink to this definition">¶</a></dt>
-<dd></dd></dl>
-
-<dl class="function">
-<dt id="pytorch_tabnet.multiclass_utils.check_unique_type">
-<code class="sig-prename descclassname">pytorch_tabnet.multiclass_utils.</code><code class="sig-name descname">check_unique_type</code><span class="sig-paren">(</span><em class="sig-param">y</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/multiclass_utils.html#check_unique_type"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.multiclass_utils.check_unique_type" title="Permalink to this definition">¶</a></dt>
-<dd></dd></dl>
-
-<dl class="function">
-<dt id="pytorch_tabnet.multiclass_utils.infer_multitask_output">
-<code class="sig-prename descclassname">pytorch_tabnet.multiclass_utils.</code><code class="sig-name descname">infer_multitask_output</code><span class="sig-paren">(</span><em class="sig-param">y_train</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/multiclass_utils.html#infer_multitask_output"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.multiclass_utils.infer_multitask_output" title="Permalink to this definition">¶</a></dt>
-<dd><p>Infer output_dim from targets
-This is for multiple tasks.</p>
-<dl class="field-list simple">
-<dt class="field-odd">Parameters</dt>
-<dd class="field-odd"><p><strong>y_train</strong> (<em>np.ndarray</em>) – Training targets</p>
-</dd>
-<dt class="field-even">Returns</dt>
-<dd class="field-even"><p><ul class="simple">
-<li><p><strong>tasks_dims</strong> (<em>list</em>) – Number of classes for output</p></li>
-<li><p><strong>tasks_labels</strong> (<em>list</em>) – List of sorted list of initial classes</p></li>
-</ul>
-</p>
-</dd>
-</dl>
-</dd></dl>
-
-<dl class="function">
-<dt id="pytorch_tabnet.multiclass_utils.infer_output_dim">
-<code class="sig-prename descclassname">pytorch_tabnet.multiclass_utils.</code><code class="sig-name descname">infer_output_dim</code><span class="sig-paren">(</span><em class="sig-param">y_train</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/multiclass_utils.html#infer_output_dim"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.multiclass_utils.infer_output_dim" title="Permalink to this definition">¶</a></dt>
-<dd><p>Infer output_dim from targets</p>
-<dl class="field-list simple">
-<dt class="field-odd">Parameters</dt>
-<dd class="field-odd"><p><strong>y_train</strong> (<em>np.array</em>) – Training targets</p>
-</dd>
-<dt class="field-even">Returns</dt>
-<dd class="field-even"><p><ul class="simple">
-<li><p><strong>output_dim</strong> (<em>int</em>) – Number of classes for output</p></li>
-<li><p><strong>train_labels</strong> (<em>list</em>) – Sorted list of initial classes</p></li>
-</ul>
-</p>
-</dd>
-</dl>
-</dd></dl>
-
-<dl class="function">
-<dt id="pytorch_tabnet.multiclass_utils.is_multilabel">
-<code class="sig-prename descclassname">pytorch_tabnet.multiclass_utils.</code><code class="sig-name descname">is_multilabel</code><span class="sig-paren">(</span><em class="sig-param">y</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/multiclass_utils.html#is_multilabel"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.multiclass_utils.is_multilabel" title="Permalink to this definition">¶</a></dt>
-<dd><p>Check if <code class="docutils literal notranslate"><span class="pre">y</span></code> is in a multilabel format.</p>
-<dl class="field-list simple">
-<dt class="field-odd">Parameters</dt>
-<dd class="field-odd"><p><strong>y</strong> (<em>numpy array of shape</em><em> [</em><em>n_samples</em><em>]</em>) – Target values.</p>
-</dd>
-<dt class="field-even">Returns</dt>
-<dd class="field-even"><p><strong>out</strong> – Return <code class="docutils literal notranslate"><span class="pre">True</span></code>, if <code class="docutils literal notranslate"><span class="pre">y</span></code> is in a multilabel format, else <code class="docutils literal notranslate"><span class="pre">`False</span></code>.</p>
-</dd>
-<dt class="field-odd">Return type</dt>
-<dd class="field-odd"><p>bool</p>
-</dd>
-</dl>
-<p class="rubric">Examples</p>
-<div class="doctest highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
-<span class="gp">&gt;&gt;&gt; </span><span class="kn">from</span> <span class="nn">sklearn.utils.multiclass</span> <span class="kn">import</span> <span class="n">is_multilabel</span>
-<span class="gp">&gt;&gt;&gt; </span><span class="n">is_multilabel</span><span class="p">([</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">])</span>
-<span class="go">False</span>
-<span class="gp">&gt;&gt;&gt; </span><span class="n">is_multilabel</span><span class="p">([[</span><span class="mi">1</span><span class="p">],</span> <span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">2</span><span class="p">],</span> <span class="p">[]])</span>
-<span class="go">False</span>
-<span class="gp">&gt;&gt;&gt; </span><span class="n">is_multilabel</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">([[</span><span class="mi">1</span><span class="p">,</span> <span class="mi">0</span><span class="p">],</span> <span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">]]))</span>
-<span class="go">True</span>
-<span class="gp">&gt;&gt;&gt; </span><span class="n">is_multilabel</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">([[</span><span class="mi">1</span><span class="p">],</span> <span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="p">[</span><span class="mi">0</span><span class="p">]]))</span>
-<span class="go">False</span>
-<span class="gp">&gt;&gt;&gt; </span><span class="n">is_multilabel</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">([[</span><span class="mi">1</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">]]))</span>
-<span class="go">True</span>
-</pre></div>
-</div>
-</dd></dl>
-
-<dl class="function">
-<dt id="pytorch_tabnet.multiclass_utils.type_of_target">
-<code class="sig-prename descclassname">pytorch_tabnet.multiclass_utils.</code><code class="sig-name descname">type_of_target</code><span class="sig-paren">(</span><em class="sig-param">y</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/multiclass_utils.html#type_of_target"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.multiclass_utils.type_of_target" title="Permalink to this definition">¶</a></dt>
-<dd><p>Determine the type of data indicated by the target.</p>
-<p>Note that this type is the most specific type that can be inferred.
-For example:</p>
-<blockquote>
-<div><ul class="simple">
-<li><p><code class="docutils literal notranslate"><span class="pre">binary</span></code> is more specific but compatible with <code class="docutils literal notranslate"><span class="pre">multiclass</span></code>.</p></li>
-<li><p><code class="docutils literal notranslate"><span class="pre">multiclass</span></code> of integers is more specific but compatible with
-<code class="docutils literal notranslate"><span class="pre">continuous</span></code>.</p></li>
-<li><p><code class="docutils literal notranslate"><span class="pre">multilabel-indicator</span></code> is more specific but compatible with
-<code class="docutils literal notranslate"><span class="pre">multiclass-multioutput</span></code>.</p></li>
-</ul>
-</div></blockquote>
-<dl class="field-list simple">
-<dt class="field-odd">Parameters</dt>
-<dd class="field-odd"><p><strong>y</strong> (<em>array-like</em>) – </p>
+<dd class="field-odd"><p><strong>weights</strong> (<em>bool</em><em> or </em><em>dict</em>) – Given weights for balancing training.</p>
 </dd>
 <dt class="field-even">Returns</dt>
-<dd class="field-even"><p><p><strong>target_type</strong> – One of:</p>
-<ul class="simple">
-<li><p>’continuous’: <cite>y</cite> is an array-like of floats that are not all
-integers, and is 1d or a column vector.</p></li>
-<li><p>’continuous-multioutput’: <cite>y</cite> is a 2d array of floats that are
-not all integers, and both dimensions are of size &gt; 1.</p></li>
-<li><p>’binary’: <cite>y</cite> contains &lt;= 2 discrete values and is 1d or a column
-vector.</p></li>
-<li><p>’multiclass’: <cite>y</cite> contains more than two discrete values, is not a
-sequence of sequences, and is 1d or a column vector.</p></li>
-<li><p>’multiclass-multioutput’: <cite>y</cite> is a 2d array that contains more
-than two discrete values, is not a sequence of sequences, and both
-dimensions are of size &gt; 1.</p></li>
-<li><p>’multilabel-indicator’: <cite>y</cite> is a label indicator matrix, an array
-of two dimensions with at least two columns, and at most 2 unique
-values.</p></li>
-<li><p>’unknown’: <cite>y</cite> is array-like but none of the above, such as a 3d
-array, sequence of sequences, or an array of non-sequence objects.</p></li>
-</ul>
-</p>
+<dd class="field-even"><p>Same bool if weights are bool, updated dict otherwise.</p>
 </dd>
 <dt class="field-odd">Return type</dt>
-<dd class="field-odd"><p>string</p>
+<dd class="field-odd"><p>bool or dict</p>
 </dd>
 </dl>
-<p class="rubric">Examples</p>
-<div class="doctest highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
-<span class="gp">&gt;&gt;&gt; </span><span class="n">type_of_target</span><span class="p">([</span><span class="mf">0.1</span><span class="p">,</span> <span class="mf">0.6</span><span class="p">])</span>
-<span class="go">&#39;continuous&#39;</span>
-<span class="gp">&gt;&gt;&gt; </span><span class="n">type_of_target</span><span class="p">([</span><span class="mi">1</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">])</span>
-<span class="go">&#39;binary&#39;</span>
-<span class="gp">&gt;&gt;&gt; </span><span class="n">type_of_target</span><span class="p">([</span><span class="s1">&#39;a&#39;</span><span class="p">,</span> <span class="s1">&#39;b&#39;</span><span class="p">,</span> <span class="s1">&#39;a&#39;</span><span class="p">])</span>
-<span class="go">&#39;binary&#39;</span>
-<span class="gp">&gt;&gt;&gt; </span><span class="n">type_of_target</span><span class="p">([</span><span class="mf">1.0</span><span class="p">,</span> <span class="mf">2.0</span><span class="p">])</span>
-<span class="go">&#39;binary&#39;</span>
-<span class="gp">&gt;&gt;&gt; </span><span class="n">type_of_target</span><span class="p">([</span><span class="mi">1</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">2</span><span class="p">])</span>
-<span class="go">&#39;multiclass&#39;</span>
-<span class="gp">&gt;&gt;&gt; </span><span class="n">type_of_target</span><span class="p">([</span><span class="mf">1.0</span><span class="p">,</span> <span class="mf">0.0</span><span class="p">,</span> <span class="mf">3.0</span><span class="p">])</span>
-<span class="go">&#39;multiclass&#39;</span>
-<span class="gp">&gt;&gt;&gt; </span><span class="n">type_of_target</span><span class="p">([</span><span class="s1">&#39;a&#39;</span><span class="p">,</span> <span class="s1">&#39;b&#39;</span><span class="p">,</span> <span class="s1">&#39;c&#39;</span><span class="p">])</span>
-<span class="go">&#39;multiclass&#39;</span>
-<span class="gp">&gt;&gt;&gt; </span><span class="n">type_of_target</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">([[</span><span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">],</span> <span class="p">[</span><span class="mi">3</span><span class="p">,</span> <span class="mi">1</span><span class="p">]]))</span>
-<span class="go">&#39;multiclass-multioutput&#39;</span>
-<span class="gp">&gt;&gt;&gt; </span><span class="n">type_of_target</span><span class="p">([[</span><span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">]])</span>
-<span class="go">&#39;multiclass-multioutput&#39;</span>
-<span class="gp">&gt;&gt;&gt; </span><span class="n">type_of_target</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">([[</span><span class="mf">1.5</span><span class="p">,</span> <span class="mf">2.0</span><span class="p">],</span> <span class="p">[</span><span class="mf">3.0</span><span class="p">,</span> <span class="mf">1.6</span><span class="p">]]))</span>
-<span class="go">&#39;continuous-multioutput&#39;</span>
-<span class="gp">&gt;&gt;&gt; </span><span class="n">type_of_target</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">([[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">],</span> <span class="p">[</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">]]))</span>
-<span class="go">&#39;multilabel-indicator&#39;</span>
-</pre></div>
-</div>
 </dd></dl>
 
-<dl class="function">
-<dt id="pytorch_tabnet.multiclass_utils.unique_labels">
-<code class="sig-prename descclassname">pytorch_tabnet.multiclass_utils.</code><code class="sig-name descname">unique_labels</code><span class="sig-paren">(</span><em class="sig-param">*ys</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/multiclass_utils.html#unique_labels"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.multiclass_utils.unique_labels" title="Permalink to this definition">¶</a></dt>
-<dd><p>Extract an ordered array of unique labels</p>
-<dl class="simple">
-<dt>We don’t allow:</dt><dd><ul class="simple">
-<li><p>mix of multilabel and multiclass (single label) targets</p></li>
-<li><p>mix of label indicator matrix and anything else,
-because there are no explicit labels)</p></li>
-<li><p>mix of label indicator matrices of different sizes</p></li>
-<li><p>mix of string and integer labels</p></li>
-</ul>
-</dd>
-</dl>
-<p>At the moment, we also don’t allow “multiclass-multioutput” input type.</p>
-<dl class="field-list simple">
-<dt class="field-odd">Parameters</dt>
-<dd class="field-odd"><p><strong>*ys</strong> (<em>array-likes</em>) – </p>
-</dd>
-<dt class="field-even">Returns</dt>
-<dd class="field-even"><p><strong>out</strong> – An ordered array of unique labels.</p>
-</dd>
-<dt class="field-odd">Return type</dt>
-<dd class="field-odd"><p>numpy array of shape [n_unique_labels]</p>
-</dd>
-</dl>
-<p class="rubric">Examples</p>
-<div class="doctest highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="kn">from</span> <span class="nn">sklearn.utils.multiclass</span> <span class="kn">import</span> <span class="n">unique_labels</span>
-<span class="gp">&gt;&gt;&gt; </span><span class="n">unique_labels</span><span class="p">([</span><span class="mi">3</span><span class="p">,</span> <span class="mi">5</span><span class="p">,</span> <span class="mi">5</span><span class="p">,</span> <span class="mi">5</span><span class="p">,</span> <span class="mi">7</span><span class="p">,</span> <span class="mi">7</span><span class="p">])</span>
-<span class="go">array([3, 5, 7])</span>
-<span class="gp">&gt;&gt;&gt; </span><span class="n">unique_labels</span><span class="p">([</span><span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="mi">4</span><span class="p">],</span> <span class="p">[</span><span class="mi">2</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="mi">4</span><span class="p">])</span>
-<span class="go">array([1, 2, 3, 4])</span>
-<span class="gp">&gt;&gt;&gt; </span><span class="n">unique_labels</span><span class="p">([</span><span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">10</span><span class="p">],</span> <span class="p">[</span><span class="mi">5</span><span class="p">,</span> <span class="mi">11</span><span class="p">])</span>
-<span class="go">array([ 1,  2,  5, 10, 11])</span>
-</pre></div>
-</div>
 </dd></dl>
 
-</div>
-</div>
-<div class="section" id="module-pytorch_tabnet.tab_model">
-<span id="pytorch-tabnet-tab-model-module"></span><h2>pytorch_tabnet.tab_model module<a class="headerlink" href="#module-pytorch_tabnet.tab_model" title="Permalink to this headline">¶</a></h2>
 <dl class="class">
-<dt id="pytorch_tabnet.tab_model.TabNetClassifier">
-<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.tab_model.</code><code class="sig-name descname">TabNetClassifier</code><span class="sig-paren">(</span><em class="sig-param">n_d: int = 8</em>, <em class="sig-param">n_a: int = 8</em>, <em class="sig-param">n_steps: int = 3</em>, <em class="sig-param">gamma: float = 1.3</em>, <em class="sig-param">cat_idxs: List[int] = &lt;factory&gt;</em>, <em class="sig-param">cat_dims: List[int] = &lt;factory&gt;</em>, <em class="sig-param">cat_emb_dim: int = 1</em>, <em class="sig-param">n_independent: int = 2</em>, <em class="sig-param">n_shared: int = 2</em>, <em class="sig-param">epsilon: float = 1e-15</em>, <em class="sig-param">momentum: float = 0.02</em>, <em class="sig-param">lambda_sparse: float = 0.001</em>, <em class="sig-param">seed: int = 0</em>, <em class="sig-param">clip_value: int = 1</em>, <em class="sig-param">verbose: int = 1</em>, <em class="sig-param">optimizer_fn: Any = &lt;class 'torch.optim.adam.Adam'&gt;</em>, <em class="sig-param">optimizer_params: Dict = &lt;factory&gt;</em>, <em class="sig-param">scheduler_fn: Any = None</em>, <em class="sig-param">scheduler_params: Dict = &lt;factory&gt;</em>, <em class="sig-param">mask_type: str = 'sparsemax'</em>, <em class="sig-param">input_dim: int = None</em>, <em class="sig-param">output_dim: int = None</em>, <em class="sig-param">device_name: str = 'auto'</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/tab_model.html#TabNetClassifier"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetClassifier" title="Permalink to this definition">¶</a></dt>
+<dt id="pytorch_tabnet.tab_model.TabNetRegressor">
+<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.tab_model.</code><code class="sig-name descname">TabNetRegressor</code><span class="sig-paren">(</span><em class="sig-param">n_d: int = 8</em>, <em class="sig-param">n_a: int = 8</em>, <em class="sig-param">n_steps: int = 3</em>, <em class="sig-param">gamma: float = 1.3</em>, <em class="sig-param">cat_idxs: List[int] = &lt;factory&gt;</em>, <em class="sig-param">cat_dims: List[int] = &lt;factory&gt;</em>, <em class="sig-param">cat_emb_dim: int = 1</em>, <em class="sig-param">n_independent: int = 2</em>, <em class="sig-param">n_shared: int = 2</em>, <em class="sig-param">epsilon: float = 1e-15</em>, <em class="sig-param">momentum: float = 0.02</em>, <em class="sig-param">lambda_sparse: float = 0.001</em>, <em class="sig-param">seed: int = 0</em>, <em class="sig-param">clip_value: int = 1</em>, <em class="sig-param">verbose: int = 1</em>, <em class="sig-param">optimizer_fn: Any = &lt;class 'torch.optim.adam.Adam'&gt;</em>, <em class="sig-param">optimizer_params: Dict = &lt;factory&gt;</em>, <em class="sig-param">scheduler_fn: Any = None</em>, <em class="sig-param">scheduler_params: Dict = &lt;factory&gt;</em>, <em class="sig-param">mask_type: str = 'sparsemax'</em>, <em class="sig-param">input_dim: int = None</em>, <em class="sig-param">output_dim: int = None</em>, <em class="sig-param">device_name: str = 'auto'</em>, <em class="sig-param">n_shared_decoder: int = 1</em>, <em class="sig-param">n_indep_decoder: int = 1</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/tab_model.html#TabNetRegressor"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetRegressor" title="Permalink to this definition">¶</a></dt>
 <dd><p>Bases: <a class="reference internal" href="#pytorch_tabnet.abstract_model.TabModel" title="pytorch_tabnet.abstract_model.TabModel"><code class="xref py py-class docutils literal notranslate"><span class="pre">pytorch_tabnet.abstract_model.TabModel</span></code></a></p>
 <dl class="attribute">
-<dt id="pytorch_tabnet.tab_model.TabNetClassifier.cat_dims">
-<code class="sig-name descname">cat_dims</code><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetClassifier.cat_dims" title="Permalink to this definition">¶</a></dt>
+<dt id="pytorch_tabnet.tab_model.TabNetRegressor.cat_dims">
+<code class="sig-name descname">cat_dims</code><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetRegressor.cat_dims" title="Permalink to this definition">¶</a></dt>
 <dd></dd></dl>
 
 <dl class="attribute">
-<dt id="pytorch_tabnet.tab_model.TabNetClassifier.cat_idxs">
-<code class="sig-name descname">cat_idxs</code><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetClassifier.cat_idxs" title="Permalink to this definition">¶</a></dt>
+<dt id="pytorch_tabnet.tab_model.TabNetRegressor.cat_idxs">
+<code class="sig-name descname">cat_idxs</code><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetRegressor.cat_idxs" title="Permalink to this definition">¶</a></dt>
 <dd></dd></dl>
 
 <dl class="method">
-<dt id="pytorch_tabnet.tab_model.TabNetClassifier.compute_loss">
-<code class="sig-name descname">compute_loss</code><span class="sig-paren">(</span><em class="sig-param">y_pred</em>, <em class="sig-param">y_true</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/tab_model.html#TabNetClassifier.compute_loss"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetClassifier.compute_loss" title="Permalink to this definition">¶</a></dt>
+<dt id="pytorch_tabnet.tab_model.TabNetRegressor.compute_loss">
+<code class="sig-name descname">compute_loss</code><span class="sig-paren">(</span><em class="sig-param">y_pred</em>, <em class="sig-param">y_true</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/tab_model.html#TabNetRegressor.compute_loss"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetRegressor.compute_loss" title="Permalink to this definition">¶</a></dt>
 <dd><p>Compute the loss.</p>
 <dl class="field-list simple">
 <dt class="field-odd">Parameters</dt>
@@ -1088,35 +895,18 @@ <h3>Multi-class / multi-label utility function<a class="headerlink" href="#multi
 </dd></dl>
 
 <dl class="attribute">
-<dt id="pytorch_tabnet.tab_model.TabNetClassifier.optimizer_params">
-<code class="sig-name descname">optimizer_params</code><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetClassifier.optimizer_params" title="Permalink to this definition">¶</a></dt>
+<dt id="pytorch_tabnet.tab_model.TabNetRegressor.optimizer_params">
+<code class="sig-name descname">optimizer_params</code><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetRegressor.optimizer_params" title="Permalink to this definition">¶</a></dt>
 <dd></dd></dl>
 
 <dl class="method">
-<dt id="pytorch_tabnet.tab_model.TabNetClassifier.predict_func">
-<code class="sig-name descname">predict_func</code><span class="sig-paren">(</span><em class="sig-param">outputs</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/tab_model.html#TabNetClassifier.predict_func"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetClassifier.predict_func" title="Permalink to this definition">¶</a></dt>
+<dt id="pytorch_tabnet.tab_model.TabNetRegressor.predict_func">
+<code class="sig-name descname">predict_func</code><span class="sig-paren">(</span><em class="sig-param">outputs</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/tab_model.html#TabNetRegressor.predict_func"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetRegressor.predict_func" title="Permalink to this definition">¶</a></dt>
 <dd></dd></dl>
 
 <dl class="method">
-<dt id="pytorch_tabnet.tab_model.TabNetClassifier.predict_proba">
-<code class="sig-name descname">predict_proba</code><span class="sig-paren">(</span><em class="sig-param">X</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/tab_model.html#TabNetClassifier.predict_proba"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetClassifier.predict_proba" title="Permalink to this definition">¶</a></dt>
-<dd><p>Make predictions for classification on a batch (valid)</p>
-<dl class="field-list simple">
-<dt class="field-odd">Parameters</dt>
-<dd class="field-odd"><p><strong>X</strong> (a :tensor: <cite>torch.Tensor</cite>) – Input data</p>
-</dd>
-<dt class="field-even">Returns</dt>
-<dd class="field-even"><p><strong>res</strong></p>
-</dd>
-<dt class="field-odd">Return type</dt>
-<dd class="field-odd"><p>np.ndarray</p>
-</dd>
-</dl>
-</dd></dl>
-
-<dl class="method">
-<dt id="pytorch_tabnet.tab_model.TabNetClassifier.prepare_target">
-<code class="sig-name descname">prepare_target</code><span class="sig-paren">(</span><em class="sig-param">y</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/tab_model.html#TabNetClassifier.prepare_target"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetClassifier.prepare_target" title="Permalink to this definition">¶</a></dt>
+<dt id="pytorch_tabnet.tab_model.TabNetRegressor.prepare_target">
+<code class="sig-name descname">prepare_target</code><span class="sig-paren">(</span><em class="sig-param">y</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/tab_model.html#TabNetRegressor.prepare_target"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetRegressor.prepare_target" title="Permalink to this definition">¶</a></dt>
 <dd><p>Prepare target before training.</p>
 <dl class="field-list simple">
 <dt class="field-odd">Parameters</dt>
@@ -1132,19 +922,19 @@ <h3>Multi-class / multi-label utility function<a class="headerlink" href="#multi
 </dd></dl>
 
 <dl class="attribute">
-<dt id="pytorch_tabnet.tab_model.TabNetClassifier.scheduler_params">
-<code class="sig-name descname">scheduler_params</code><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetClassifier.scheduler_params" title="Permalink to this definition">¶</a></dt>
+<dt id="pytorch_tabnet.tab_model.TabNetRegressor.scheduler_params">
+<code class="sig-name descname">scheduler_params</code><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetRegressor.scheduler_params" title="Permalink to this definition">¶</a></dt>
 <dd></dd></dl>
 
 <dl class="method">
-<dt id="pytorch_tabnet.tab_model.TabNetClassifier.stack_batches">
-<code class="sig-name descname">stack_batches</code><span class="sig-paren">(</span><em class="sig-param">list_y_true</em>, <em class="sig-param">list_y_score</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/tab_model.html#TabNetClassifier.stack_batches"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetClassifier.stack_batches" title="Permalink to this definition">¶</a></dt>
+<dt id="pytorch_tabnet.tab_model.TabNetRegressor.stack_batches">
+<code class="sig-name descname">stack_batches</code><span class="sig-paren">(</span><em class="sig-param">list_y_true</em>, <em class="sig-param">list_y_score</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/tab_model.html#TabNetRegressor.stack_batches"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetRegressor.stack_batches" title="Permalink to this definition">¶</a></dt>
 <dd></dd></dl>
 
 <dl class="method">
-<dt id="pytorch_tabnet.tab_model.TabNetClassifier.update_fit_params">
-<code class="sig-name descname">update_fit_params</code><span class="sig-paren">(</span><em class="sig-param">X_train</em>, <em class="sig-param">y_train</em>, <em class="sig-param">eval_set</em>, <em class="sig-param">weights</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/tab_model.html#TabNetClassifier.update_fit_params"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetClassifier.update_fit_params" title="Permalink to this definition">¶</a></dt>
-<dd><p>Set attributes relative to fit function.</p>
+<dt id="pytorch_tabnet.tab_model.TabNetRegressor.update_fit_params">
+<code class="sig-name descname">update_fit_params</code><span class="sig-paren">(</span><em class="sig-param">X_train</em>, <em class="sig-param">y_train</em>, <em class="sig-param">eval_set</em>, <em class="sig-param">weights</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/tab_model.html#TabNetRegressor.update_fit_params"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetRegressor.update_fit_params" title="Permalink to this definition">¶</a></dt>
+<dd><p>Set attributes relative to fit function.</p>
 <dl class="field-list simple">
 <dt class="field-odd">Parameters</dt>
 <dd class="field-odd"><ul class="simple">
@@ -1158,404 +948,190 @@ <h3>Multi-class / multi-label utility function<a class="headerlink" href="#multi
 </dl>
 </dd></dl>
 
-<dl class="method">
-<dt id="pytorch_tabnet.tab_model.TabNetClassifier.weight_updater">
-<code class="sig-name descname">weight_updater</code><span class="sig-paren">(</span><em class="sig-param">weights</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/tab_model.html#TabNetClassifier.weight_updater"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetClassifier.weight_updater" title="Permalink to this definition">¶</a></dt>
-<dd><p>Updates weights dictionary according to target_mapper.</p>
-<dl class="field-list simple">
-<dt class="field-odd">Parameters</dt>
-<dd class="field-odd"><p><strong>weights</strong> (<em>bool</em><em> or </em><em>dict</em>) – Given weights for balancing training.</p>
-</dd>
-<dt class="field-even">Returns</dt>
-<dd class="field-even"><p>Same bool if weights are bool, updated dict otherwise.</p>
-</dd>
-<dt class="field-odd">Return type</dt>
-<dd class="field-odd"><p>bool or dict</p>
-</dd>
-</dl>
-</dd></dl>
-
 </dd></dl>
 
+</section>
+<section id="module-pytorch_tabnet.sparsemax">
+<span id="pytorch-tabnet-sparsemax-module"></span><h2>pytorch_tabnet.sparsemax module<a class="headerlink" href="#module-pytorch_tabnet.sparsemax" title="Permalink to this headline">¶</a></h2>
 <dl class="class">
-<dt id="pytorch_tabnet.tab_model.TabNetRegressor">
-<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.tab_model.</code><code class="sig-name descname">TabNetRegressor</code><span class="sig-paren">(</span><em class="sig-param">n_d: int = 8</em>, <em class="sig-param">n_a: int = 8</em>, <em class="sig-param">n_steps: int = 3</em>, <em class="sig-param">gamma: float = 1.3</em>, <em class="sig-param">cat_idxs: List[int] = &lt;factory&gt;</em>, <em class="sig-param">cat_dims: List[int] = &lt;factory&gt;</em>, <em class="sig-param">cat_emb_dim: int = 1</em>, <em class="sig-param">n_independent: int = 2</em>, <em class="sig-param">n_shared: int = 2</em>, <em class="sig-param">epsilon: float = 1e-15</em>, <em class="sig-param">momentum: float = 0.02</em>, <em class="sig-param">lambda_sparse: float = 0.001</em>, <em class="sig-param">seed: int = 0</em>, <em class="sig-param">clip_value: int = 1</em>, <em class="sig-param">verbose: int = 1</em>, <em class="sig-param">optimizer_fn: Any = &lt;class 'torch.optim.adam.Adam'&gt;</em>, <em class="sig-param">optimizer_params: Dict = &lt;factory&gt;</em>, <em class="sig-param">scheduler_fn: Any = None</em>, <em class="sig-param">scheduler_params: Dict = &lt;factory&gt;</em>, <em class="sig-param">mask_type: str = 'sparsemax'</em>, <em class="sig-param">input_dim: int = None</em>, <em class="sig-param">output_dim: int = None</em>, <em class="sig-param">device_name: str = 'auto'</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/tab_model.html#TabNetRegressor"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetRegressor" title="Permalink to this definition">¶</a></dt>
-<dd><p>Bases: <a class="reference internal" href="#pytorch_tabnet.abstract_model.TabModel" title="pytorch_tabnet.abstract_model.TabModel"><code class="xref py py-class docutils literal notranslate"><span class="pre">pytorch_tabnet.abstract_model.TabModel</span></code></a></p>
-<dl class="attribute">
-<dt id="pytorch_tabnet.tab_model.TabNetRegressor.cat_dims">
-<code class="sig-name descname">cat_dims</code><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetRegressor.cat_dims" title="Permalink to this definition">¶</a></dt>
-<dd></dd></dl>
-
-<dl class="attribute">
-<dt id="pytorch_tabnet.tab_model.TabNetRegressor.cat_idxs">
-<code class="sig-name descname">cat_idxs</code><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetRegressor.cat_idxs" title="Permalink to this definition">¶</a></dt>
-<dd></dd></dl>
-
+<dt id="pytorch_tabnet.sparsemax.Entmax15">
+<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.sparsemax.</code><code class="sig-name descname">Entmax15</code><span class="sig-paren">(</span><em class="sig-param">dim=-1</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/sparsemax.html#Entmax15"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.sparsemax.Entmax15" title="Permalink to this definition">¶</a></dt>
+<dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">torch.nn.modules.module.Module</span></code></p>
 <dl class="method">
-<dt id="pytorch_tabnet.tab_model.TabNetRegressor.compute_loss">
-<code class="sig-name descname">compute_loss</code><span class="sig-paren">(</span><em class="sig-param">y_pred</em>, <em class="sig-param">y_true</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/tab_model.html#TabNetRegressor.compute_loss"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetRegressor.compute_loss" title="Permalink to this definition">¶</a></dt>
-<dd><p>Compute the loss.</p>
-<dl class="field-list simple">
-<dt class="field-odd">Parameters</dt>
-<dd class="field-odd"><ul class="simple">
-<li><p><strong>y_score</strong> (a :tensor: <cite>torch.Tensor</cite>) – Score matrix</p></li>
-<li><p><strong>y_true</strong> (a :tensor: <cite>torch.Tensor</cite>) – Target matrix</p></li>
-</ul>
-</dd>
-<dt class="field-even">Returns</dt>
-<dd class="field-even"><p>Loss value</p>
-</dd>
-<dt class="field-odd">Return type</dt>
-<dd class="field-odd"><p>float</p>
-</dd>
-</dl>
+<dt id="pytorch_tabnet.sparsemax.Entmax15.forward">
+<code class="sig-name descname">forward</code><span class="sig-paren">(</span><em class="sig-param">input</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/sparsemax.html#Entmax15.forward"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.sparsemax.Entmax15.forward" title="Permalink to this definition">¶</a></dt>
+<dd><p>Defines the computation performed at every call.</p>
+<p>Should be overridden by all subclasses.</p>
+<div class="admonition note">
+<p class="admonition-title">Note</p>
+<p>Although the recipe for forward pass needs to be defined within
+this function, one should call the <code class="xref py py-class docutils literal notranslate"><span class="pre">Module</span></code> instance afterwards
+instead of this since the former takes care of running the
+registered hooks while the latter silently ignores them.</p>
+</div>
 </dd></dl>
 
-<dl class="attribute">
-<dt id="pytorch_tabnet.tab_model.TabNetRegressor.optimizer_params">
-<code class="sig-name descname">optimizer_params</code><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetRegressor.optimizer_params" title="Permalink to this definition">¶</a></dt>
-<dd></dd></dl>
+</dd></dl>
 
+<dl class="class">
+<dt id="pytorch_tabnet.sparsemax.Entmax15Function">
+<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.sparsemax.</code><code class="sig-name descname">Entmax15Function</code><a class="reference internal" href="../_modules/pytorch_tabnet/sparsemax.html#Entmax15Function"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.sparsemax.Entmax15Function" title="Permalink to this definition">¶</a></dt>
+<dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">torch.autograd.function.Function</span></code></p>
+<p>An implementation of exact Entmax with alpha=1.5 (B. Peters, V. Niculae, A. Martins). See
+:cite:<a href="#id1"><span class="problematic" id="id2">`</span></a><a class="reference external" href="https://arxiv.org/abs/1905.05702">https://arxiv.org/abs/1905.05702</a> for detailed description.
+Source: <a class="reference external" href="https://github.com/deep-spin/entmax">https://github.com/deep-spin/entmax</a></p>
 <dl class="method">
-<dt id="pytorch_tabnet.tab_model.TabNetRegressor.predict_func">
-<code class="sig-name descname">predict_func</code><span class="sig-paren">(</span><em class="sig-param">outputs</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/tab_model.html#TabNetRegressor.predict_func"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetRegressor.predict_func" title="Permalink to this definition">¶</a></dt>
-<dd></dd></dl>
+<dt id="pytorch_tabnet.sparsemax.Entmax15Function.backward">
+<em class="property">static </em><code class="sig-name descname">backward</code><span class="sig-paren">(</span><em class="sig-param">ctx</em>, <em class="sig-param">grad_output</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/sparsemax.html#Entmax15Function.backward"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.sparsemax.Entmax15Function.backward" title="Permalink to this definition">¶</a></dt>
+<dd><p>Defines a formula for differentiating the operation.</p>
+<p>This function is to be overridden by all subclasses.</p>
+<p>It must accept a context <code class="xref py py-attr docutils literal notranslate"><span class="pre">ctx</span></code> as the first argument, followed by
+as many outputs did <a class="reference internal" href="#pytorch_tabnet.sparsemax.Entmax15Function.forward" title="pytorch_tabnet.sparsemax.Entmax15Function.forward"><code class="xref py py-func docutils literal notranslate"><span class="pre">forward()</span></code></a> return, and it should return as many
+tensors, as there were inputs to <a class="reference internal" href="#pytorch_tabnet.sparsemax.Entmax15Function.forward" title="pytorch_tabnet.sparsemax.Entmax15Function.forward"><code class="xref py py-func docutils literal notranslate"><span class="pre">forward()</span></code></a>. Each argument is the
+gradient w.r.t the given output, and each returned value should be the
+gradient w.r.t. the corresponding input.</p>
+<p>The context can be used to retrieve tensors saved during the forward
+pass. It also has an attribute <code class="xref py py-attr docutils literal notranslate"><span class="pre">ctx.needs_input_grad</span></code> as a tuple
+of booleans representing whether each input needs gradient. E.g.,
+<a class="reference internal" href="#pytorch_tabnet.sparsemax.Entmax15Function.backward" title="pytorch_tabnet.sparsemax.Entmax15Function.backward"><code class="xref py py-func docutils literal notranslate"><span class="pre">backward()</span></code></a> will have <code class="docutils literal notranslate"><span class="pre">ctx.needs_input_grad[0]</span> <span class="pre">=</span> <span class="pre">True</span></code> if the
+first input to <a class="reference internal" href="#pytorch_tabnet.sparsemax.Entmax15Function.forward" title="pytorch_tabnet.sparsemax.Entmax15Function.forward"><code class="xref py py-func docutils literal notranslate"><span class="pre">forward()</span></code></a> needs gradient computated w.r.t. the
+output.</p>
+</dd></dl>
 
 <dl class="method">
-<dt id="pytorch_tabnet.tab_model.TabNetRegressor.prepare_target">
-<code class="sig-name descname">prepare_target</code><span class="sig-paren">(</span><em class="sig-param">y</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/tab_model.html#TabNetRegressor.prepare_target"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetRegressor.prepare_target" title="Permalink to this definition">¶</a></dt>
-<dd><p>Prepare target before training.</p>
-<dl class="field-list simple">
-<dt class="field-odd">Parameters</dt>
-<dd class="field-odd"><p><strong>y</strong> (a :tensor: <cite>torch.Tensor</cite>) – Target matrix.</p>
-</dd>
-<dt class="field-even">Returns</dt>
-<dd class="field-even"><p>Converted target matrix.</p>
-</dd>
-<dt class="field-odd">Return type</dt>
-<dd class="field-odd"><p><cite>torch.Tensor</cite></p>
-</dd>
-</dl>
+<dt id="pytorch_tabnet.sparsemax.Entmax15Function.forward">
+<em class="property">static </em><code class="sig-name descname">forward</code><span class="sig-paren">(</span><em class="sig-param">ctx</em>, <em class="sig-param">input</em>, <em class="sig-param">dim=-1</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/sparsemax.html#Entmax15Function.forward"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.sparsemax.Entmax15Function.forward" title="Permalink to this definition">¶</a></dt>
+<dd><p>Performs the operation.</p>
+<p>This function is to be overridden by all subclasses.</p>
+<p>It must accept a context ctx as the first argument, followed by any
+number of arguments (tensors or other types).</p>
+<p>The context can be used to store tensors that can be then retrieved
+during the backward pass.</p>
 </dd></dl>
 
-<dl class="attribute">
-<dt id="pytorch_tabnet.tab_model.TabNetRegressor.scheduler_params">
-<code class="sig-name descname">scheduler_params</code><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetRegressor.scheduler_params" title="Permalink to this definition">¶</a></dt>
-<dd></dd></dl>
+</dd></dl>
 
+<dl class="class">
+<dt id="pytorch_tabnet.sparsemax.Entmoid15">
+<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.sparsemax.</code><code class="sig-name descname">Entmoid15</code><a class="reference internal" href="../_modules/pytorch_tabnet/sparsemax.html#Entmoid15"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.sparsemax.Entmoid15" title="Permalink to this definition">¶</a></dt>
+<dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">torch.autograd.function.Function</span></code></p>
+<p>A highly optimized equivalent of lambda x: Entmax15([x, 0])</p>
 <dl class="method">
-<dt id="pytorch_tabnet.tab_model.TabNetRegressor.stack_batches">
-<code class="sig-name descname">stack_batches</code><span class="sig-paren">(</span><em class="sig-param">list_y_true</em>, <em class="sig-param">list_y_score</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/tab_model.html#TabNetRegressor.stack_batches"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetRegressor.stack_batches" title="Permalink to this definition">¶</a></dt>
-<dd></dd></dl>
+<dt id="pytorch_tabnet.sparsemax.Entmoid15.backward">
+<em class="property">static </em><code class="sig-name descname">backward</code><span class="sig-paren">(</span><em class="sig-param">ctx</em>, <em class="sig-param">grad_output</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/sparsemax.html#Entmoid15.backward"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.sparsemax.Entmoid15.backward" title="Permalink to this definition">¶</a></dt>
+<dd><p>Defines a formula for differentiating the operation.</p>
+<p>This function is to be overridden by all subclasses.</p>
+<p>It must accept a context <code class="xref py py-attr docutils literal notranslate"><span class="pre">ctx</span></code> as the first argument, followed by
+as many outputs did <a class="reference internal" href="#pytorch_tabnet.sparsemax.Entmoid15.forward" title="pytorch_tabnet.sparsemax.Entmoid15.forward"><code class="xref py py-func docutils literal notranslate"><span class="pre">forward()</span></code></a> return, and it should return as many
+tensors, as there were inputs to <a class="reference internal" href="#pytorch_tabnet.sparsemax.Entmoid15.forward" title="pytorch_tabnet.sparsemax.Entmoid15.forward"><code class="xref py py-func docutils literal notranslate"><span class="pre">forward()</span></code></a>. Each argument is the
+gradient w.r.t the given output, and each returned value should be the
+gradient w.r.t. the corresponding input.</p>
+<p>The context can be used to retrieve tensors saved during the forward
+pass. It also has an attribute <code class="xref py py-attr docutils literal notranslate"><span class="pre">ctx.needs_input_grad</span></code> as a tuple
+of booleans representing whether each input needs gradient. E.g.,
+<a class="reference internal" href="#pytorch_tabnet.sparsemax.Entmoid15.backward" title="pytorch_tabnet.sparsemax.Entmoid15.backward"><code class="xref py py-func docutils literal notranslate"><span class="pre">backward()</span></code></a> will have <code class="docutils literal notranslate"><span class="pre">ctx.needs_input_grad[0]</span> <span class="pre">=</span> <span class="pre">True</span></code> if the
+first input to <a class="reference internal" href="#pytorch_tabnet.sparsemax.Entmoid15.forward" title="pytorch_tabnet.sparsemax.Entmoid15.forward"><code class="xref py py-func docutils literal notranslate"><span class="pre">forward()</span></code></a> needs gradient computated w.r.t. the
+output.</p>
+</dd></dl>
 
 <dl class="method">
-<dt id="pytorch_tabnet.tab_model.TabNetRegressor.update_fit_params">
-<code class="sig-name descname">update_fit_params</code><span class="sig-paren">(</span><em class="sig-param">X_train</em>, <em class="sig-param">y_train</em>, <em class="sig-param">eval_set</em>, <em class="sig-param">weights</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/tab_model.html#TabNetRegressor.update_fit_params"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.tab_model.TabNetRegressor.update_fit_params" title="Permalink to this definition">¶</a></dt>
-<dd><p>Set attributes relative to fit function.</p>
-<dl class="field-list simple">
-<dt class="field-odd">Parameters</dt>
-<dd class="field-odd"><ul class="simple">
-<li><p><strong>X_train</strong> (<em>np.ndarray</em>) – Train set</p></li>
-<li><p><strong>y_train</strong> (<em>np.array</em>) – Train targets</p></li>
-<li><p><strong>eval_set</strong> (<em>list of tuple</em>) – List of eval tuple set (X, y).</p></li>
-<li><p><strong>weights</strong> (<em>bool</em><em> or </em><em>dictionnary</em>) – 0 for no balancing
-1 for automated balancing</p></li>
-</ul>
-</dd>
-</dl>
+<dt id="pytorch_tabnet.sparsemax.Entmoid15.forward">
+<em class="property">static </em><code class="sig-name descname">forward</code><span class="sig-paren">(</span><em class="sig-param">ctx</em>, <em class="sig-param">input</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/sparsemax.html#Entmoid15.forward"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.sparsemax.Entmoid15.forward" title="Permalink to this definition">¶</a></dt>
+<dd><p>Performs the operation.</p>
+<p>This function is to be overridden by all subclasses.</p>
+<p>It must accept a context ctx as the first argument, followed by any
+number of arguments (tensors or other types).</p>
+<p>The context can be used to store tensors that can be then retrieved
+during the backward pass.</p>
 </dd></dl>
 
 </dd></dl>
 
-</div>
-<div class="section" id="module-pytorch_tabnet.abstract_model">
-<span id="pytorch-tabnet-abstract-model-module"></span><h2>pytorch_tabnet.abstract_model module<a class="headerlink" href="#module-pytorch_tabnet.abstract_model" title="Permalink to this headline">¶</a></h2>
 <dl class="class">
-<dt id="pytorch_tabnet.abstract_model.TabModel">
-<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.abstract_model.</code><code class="sig-name descname">TabModel</code><span class="sig-paren">(</span><em class="sig-param">n_d: int = 8</em>, <em class="sig-param">n_a: int = 8</em>, <em class="sig-param">n_steps: int = 3</em>, <em class="sig-param">gamma: float = 1.3</em>, <em class="sig-param">cat_idxs: List[int] = &lt;factory&gt;</em>, <em class="sig-param">cat_dims: List[int] = &lt;factory&gt;</em>, <em class="sig-param">cat_emb_dim: int = 1</em>, <em class="sig-param">n_independent: int = 2</em>, <em class="sig-param">n_shared: int = 2</em>, <em class="sig-param">epsilon: float = 1e-15</em>, <em class="sig-param">momentum: float = 0.02</em>, <em class="sig-param">lambda_sparse: float = 0.001</em>, <em class="sig-param">seed: int = 0</em>, <em class="sig-param">clip_value: int = 1</em>, <em class="sig-param">verbose: int = 1</em>, <em class="sig-param">optimizer_fn: Any = &lt;class 'torch.optim.adam.Adam'&gt;</em>, <em class="sig-param">optimizer_params: Dict = &lt;factory&gt;</em>, <em class="sig-param">scheduler_fn: Any = None</em>, <em class="sig-param">scheduler_params: Dict = &lt;factory&gt;</em>, <em class="sig-param">mask_type: str = 'sparsemax'</em>, <em class="sig-param">input_dim: int = None</em>, <em class="sig-param">output_dim: int = None</em>, <em class="sig-param">device_name: str = 'auto'</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/abstract_model.html#TabModel"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel" title="Permalink to this definition">¶</a></dt>
-<dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">sklearn.base.BaseEstimator</span></code></p>
-<p>Class for TabNet model.</p>
-<dl class="attribute">
-<dt id="pytorch_tabnet.abstract_model.TabModel.cat_dims">
-<code class="sig-name descname">cat_dims</code><em class="property">: List[int]</em><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.cat_dims" title="Permalink to this definition">¶</a></dt>
-<dd></dd></dl>
-
-<dl class="attribute">
-<dt id="pytorch_tabnet.abstract_model.TabModel.cat_emb_dim">
-<code class="sig-name descname">cat_emb_dim</code><em class="property">: int</em><em class="property"> = 1</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.cat_emb_dim" title="Permalink to this definition">¶</a></dt>
-<dd></dd></dl>
+<dt id="pytorch_tabnet.sparsemax.Sparsemax">
+<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.sparsemax.</code><code class="sig-name descname">Sparsemax</code><span class="sig-paren">(</span><em class="sig-param">dim=-1</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/sparsemax.html#Sparsemax"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.sparsemax.Sparsemax" title="Permalink to this definition">¶</a></dt>
+<dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">torch.nn.modules.module.Module</span></code></p>
+<dl class="method">
+<dt id="pytorch_tabnet.sparsemax.Sparsemax.forward">
+<code class="sig-name descname">forward</code><span class="sig-paren">(</span><em class="sig-param">input</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/sparsemax.html#Sparsemax.forward"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.sparsemax.Sparsemax.forward" title="Permalink to this definition">¶</a></dt>
+<dd><p>Defines the computation performed at every call.</p>
+<p>Should be overridden by all subclasses.</p>
+<div class="admonition note">
+<p class="admonition-title">Note</p>
+<p>Although the recipe for forward pass needs to be defined within
+this function, one should call the <code class="xref py py-class docutils literal notranslate"><span class="pre">Module</span></code> instance afterwards
+instead of this since the former takes care of running the
+registered hooks while the latter silently ignores them.</p>
+</div>
+</dd></dl>
 
-<dl class="attribute">
-<dt id="pytorch_tabnet.abstract_model.TabModel.cat_idxs">
-<code class="sig-name descname">cat_idxs</code><em class="property">: List[int]</em><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.cat_idxs" title="Permalink to this definition">¶</a></dt>
-<dd></dd></dl>
+</dd></dl>
 
-<dl class="attribute">
-<dt id="pytorch_tabnet.abstract_model.TabModel.clip_value">
-<code class="sig-name descname">clip_value</code><em class="property">: int</em><em class="property"> = 1</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.clip_value" title="Permalink to this definition">¶</a></dt>
-<dd></dd></dl>
+<dl class="class">
+<dt id="pytorch_tabnet.sparsemax.SparsemaxFunction">
+<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.sparsemax.</code><code class="sig-name descname">SparsemaxFunction</code><a class="reference internal" href="../_modules/pytorch_tabnet/sparsemax.html#SparsemaxFunction"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.sparsemax.SparsemaxFunction" title="Permalink to this definition">¶</a></dt>
+<dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">torch.autograd.function.Function</span></code></p>
+<p>An implementation of sparsemax (Martins &amp; Astudillo, 2016). See
+<a href="#id3"><span class="problematic" id="id4">:cite:`DBLP:journals/corr/MartinsA16`</span></a> for detailed description.
+By Ben Peters and Vlad Niculae</p>
+<dl class="method">
+<dt id="pytorch_tabnet.sparsemax.SparsemaxFunction.backward">
+<em class="property">static </em><code class="sig-name descname">backward</code><span class="sig-paren">(</span><em class="sig-param">ctx</em>, <em class="sig-param">grad_output</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/sparsemax.html#SparsemaxFunction.backward"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.sparsemax.SparsemaxFunction.backward" title="Permalink to this definition">¶</a></dt>
+<dd><p>Defines a formula for differentiating the operation.</p>
+<p>This function is to be overridden by all subclasses.</p>
+<p>It must accept a context <code class="xref py py-attr docutils literal notranslate"><span class="pre">ctx</span></code> as the first argument, followed by
+as many outputs did <a class="reference internal" href="#pytorch_tabnet.sparsemax.SparsemaxFunction.forward" title="pytorch_tabnet.sparsemax.SparsemaxFunction.forward"><code class="xref py py-func docutils literal notranslate"><span class="pre">forward()</span></code></a> return, and it should return as many
+tensors, as there were inputs to <a class="reference internal" href="#pytorch_tabnet.sparsemax.SparsemaxFunction.forward" title="pytorch_tabnet.sparsemax.SparsemaxFunction.forward"><code class="xref py py-func docutils literal notranslate"><span class="pre">forward()</span></code></a>. Each argument is the
+gradient w.r.t the given output, and each returned value should be the
+gradient w.r.t. the corresponding input.</p>
+<p>The context can be used to retrieve tensors saved during the forward
+pass. It also has an attribute <code class="xref py py-attr docutils literal notranslate"><span class="pre">ctx.needs_input_grad</span></code> as a tuple
+of booleans representing whether each input needs gradient. E.g.,
+<a class="reference internal" href="#pytorch_tabnet.sparsemax.SparsemaxFunction.backward" title="pytorch_tabnet.sparsemax.SparsemaxFunction.backward"><code class="xref py py-func docutils literal notranslate"><span class="pre">backward()</span></code></a> will have <code class="docutils literal notranslate"><span class="pre">ctx.needs_input_grad[0]</span> <span class="pre">=</span> <span class="pre">True</span></code> if the
+first input to <a class="reference internal" href="#pytorch_tabnet.sparsemax.SparsemaxFunction.forward" title="pytorch_tabnet.sparsemax.SparsemaxFunction.forward"><code class="xref py py-func docutils literal notranslate"><span class="pre">forward()</span></code></a> needs gradient computated w.r.t. the
+output.</p>
+</dd></dl>
 
 <dl class="method">
-<dt id="pytorch_tabnet.abstract_model.TabModel.compute_loss">
-<em class="property">abstract </em><code class="sig-name descname">compute_loss</code><span class="sig-paren">(</span><em class="sig-param">y_score</em>, <em class="sig-param">y_true</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/abstract_model.html#TabModel.compute_loss"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.compute_loss" title="Permalink to this definition">¶</a></dt>
-<dd><p>Compute the loss.</p>
+<dt id="pytorch_tabnet.sparsemax.SparsemaxFunction.forward">
+<em class="property">static </em><code class="sig-name descname">forward</code><span class="sig-paren">(</span><em class="sig-param">ctx</em>, <em class="sig-param">input</em>, <em class="sig-param">dim=-1</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/sparsemax.html#SparsemaxFunction.forward"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.sparsemax.SparsemaxFunction.forward" title="Permalink to this definition">¶</a></dt>
+<dd><p>sparsemax: normalizing sparse transform (a la softmax)</p>
 <dl class="field-list simple">
 <dt class="field-odd">Parameters</dt>
 <dd class="field-odd"><ul class="simple">
-<li><p><strong>y_score</strong> (a :tensor: <cite>torch.Tensor</cite>) – Score matrix</p></li>
-<li><p><strong>y_true</strong> (a :tensor: <cite>torch.Tensor</cite>) – Target matrix</p></li>
+<li><p><strong>ctx</strong> (<em>torch.autograd.function._ContextMethodMixin</em>) – </p></li>
+<li><p><strong>input</strong> (<em>torch.Tensor</em>) – any shape</p></li>
+<li><p><strong>dim</strong> (<em>int</em>) – dimension along which to apply sparsemax</p></li>
 </ul>
 </dd>
 <dt class="field-even">Returns</dt>
-<dd class="field-even"><p>Loss value</p>
+<dd class="field-even"><p><strong>output</strong> – same shape as input</p>
 </dd>
 <dt class="field-odd">Return type</dt>
-<dd class="field-odd"><p>float</p>
+<dd class="field-odd"><p>torch.Tensor</p>
 </dd>
 </dl>
 </dd></dl>
 
-<dl class="attribute">
-<dt id="pytorch_tabnet.abstract_model.TabModel.device_name">
-<code class="sig-name descname">device_name</code><em class="property">: str</em><em class="property"> = 'auto'</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.device_name" title="Permalink to this definition">¶</a></dt>
+</dd></dl>
+
+<dl class="function">
+<dt id="pytorch_tabnet.sparsemax.entmax15">
+<code class="sig-prename descclassname">pytorch_tabnet.sparsemax.</code><code class="sig-name descname">entmax15</code><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="headerlink" href="#pytorch_tabnet.sparsemax.entmax15" title="Permalink to this definition">¶</a></dt>
 <dd></dd></dl>
 
-<dl class="attribute">
-<dt id="pytorch_tabnet.abstract_model.TabModel.epsilon">
-<code class="sig-name descname">epsilon</code><em class="property">: float</em><em class="property"> = 1e-15</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.epsilon" title="Permalink to this definition">¶</a></dt>
-<dd></dd></dl>
-
-<dl class="method">
-<dt id="pytorch_tabnet.abstract_model.TabModel.explain">
-<code class="sig-name descname">explain</code><span class="sig-paren">(</span><em class="sig-param">X</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/abstract_model.html#TabModel.explain"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.explain" title="Permalink to this definition">¶</a></dt>
-<dd><p>Return local explanation</p>
-<dl class="field-list simple">
-<dt class="field-odd">Parameters</dt>
-<dd class="field-odd"><p><strong>X</strong> (tensor: <cite>torch.Tensor</cite>) – Input data</p>
-</dd>
-<dt class="field-even">Returns</dt>
-<dd class="field-even"><p><ul class="simple">
-<li><p><strong>M_explain</strong> (<em>matrix</em>) – Importance per sample, per columns.</p></li>
-<li><p><strong>masks</strong> (<em>matrix</em>) – Sparse matrix showing attention masks used by network.</p></li>
-</ul>
-</p>
-</dd>
-</dl>
-</dd></dl>
-
-<dl class="method">
-<dt id="pytorch_tabnet.abstract_model.TabModel.fit">
-<code class="sig-name descname">fit</code><span class="sig-paren">(</span><em class="sig-param">X_train</em>, <em class="sig-param">y_train</em>, <em class="sig-param">eval_set=None</em>, <em class="sig-param">eval_name=None</em>, <em class="sig-param">eval_metric=None</em>, <em class="sig-param">loss_fn=None</em>, <em class="sig-param">weights=0</em>, <em class="sig-param">max_epochs=100</em>, <em class="sig-param">patience=10</em>, <em class="sig-param">batch_size=1024</em>, <em class="sig-param">virtual_batch_size=128</em>, <em class="sig-param">num_workers=0</em>, <em class="sig-param">drop_last=False</em>, <em class="sig-param">callbacks=None</em>, <em class="sig-param">pin_memory=True</em>, <em class="sig-param">from_unsupervised=None</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/abstract_model.html#TabModel.fit"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.fit" title="Permalink to this definition">¶</a></dt>
-<dd><p>Train a neural network stored in self.network
-Using train_dataloader for training data and
-valid_dataloader for validation.</p>
-<dl class="field-list simple">
-<dt class="field-odd">Parameters</dt>
-<dd class="field-odd"><ul class="simple">
-<li><p><strong>X_train</strong> (<em>np.ndarray</em>) – Train set</p></li>
-<li><p><strong>y_train</strong> (<em>np.array</em>) – Train targets</p></li>
-<li><p><strong>eval_set</strong> (<em>list of tuple</em>) – List of eval tuple set (X, y).
-The last one is used for early stopping</p></li>
-<li><p><strong>eval_name</strong> (<em>list of str</em>) – List of eval set names.</p></li>
-<li><p><strong>eval_metric</strong> (<em>list of str</em>) – List of evaluation metrics.
-The last metric is used for early stopping.</p></li>
-<li><p><strong>loss_fn</strong> (<em>callable</em><em> or </em><em>None</em>) – a PyTorch loss function</p></li>
-<li><p><strong>weights</strong> (<em>bool</em><em> or </em><em>dictionnary</em>) – 0 for no balancing
-1 for automated balancing
-dict for custom weights per class</p></li>
-<li><p><strong>max_epochs</strong> (<em>int</em>) – Maximum number of epochs during training</p></li>
-<li><p><strong>patience</strong> (<em>int</em>) – Number of consecutive non improving epoch before early stopping</p></li>
-<li><p><strong>batch_size</strong> (<em>int</em>) – Training batch size</p></li>
-<li><p><strong>virtual_batch_size</strong> (<em>int</em>) – Batch size for Ghost Batch Normalization (virtual_batch_size &lt; batch_size)</p></li>
-<li><p><strong>num_workers</strong> (<em>int</em>) – Number of workers used in torch.utils.data.DataLoader</p></li>
-<li><p><strong>drop_last</strong> (<em>bool</em>) – Whether to drop last batch during training</p></li>
-<li><p><strong>callbacks</strong> (<em>list of callback function</em>) – List of custom callbacks</p></li>
-<li><p><strong>pin_memory</strong> (<em>bool</em>) – Whether to set pin_memory to True or False during training</p></li>
-<li><p><strong>from_unsupervised</strong> (<em>unsupervised trained model</em>) – Use a previously self supervised model as starting weights</p></li>
-</ul>
-</dd>
-</dl>
-</dd></dl>
-
-<dl class="attribute">
-<dt id="pytorch_tabnet.abstract_model.TabModel.gamma">
-<code class="sig-name descname">gamma</code><em class="property">: float</em><em class="property"> = 1.3</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.gamma" title="Permalink to this definition">¶</a></dt>
-<dd></dd></dl>
-
-<dl class="attribute">
-<dt id="pytorch_tabnet.abstract_model.TabModel.input_dim">
-<code class="sig-name descname">input_dim</code><em class="property">: int</em><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.input_dim" title="Permalink to this definition">¶</a></dt>
-<dd></dd></dl>
-
-<dl class="attribute">
-<dt id="pytorch_tabnet.abstract_model.TabModel.lambda_sparse">
-<code class="sig-name descname">lambda_sparse</code><em class="property">: float</em><em class="property"> = 0.001</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.lambda_sparse" title="Permalink to this definition">¶</a></dt>
-<dd></dd></dl>
-
-<dl class="method">
-<dt id="pytorch_tabnet.abstract_model.TabModel.load_class_attrs">
-<code class="sig-name descname">load_class_attrs</code><span class="sig-paren">(</span><em class="sig-param">class_attrs</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/abstract_model.html#TabModel.load_class_attrs"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.load_class_attrs" title="Permalink to this definition">¶</a></dt>
-<dd></dd></dl>
-
-<dl class="method">
-<dt id="pytorch_tabnet.abstract_model.TabModel.load_model">
-<code class="sig-name descname">load_model</code><span class="sig-paren">(</span><em class="sig-param">filepath</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/abstract_model.html#TabModel.load_model"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.load_model" title="Permalink to this definition">¶</a></dt>
-<dd><p>Load TabNet model.</p>
-<dl class="field-list simple">
-<dt class="field-odd">Parameters</dt>
-<dd class="field-odd"><p><strong>filepath</strong> (<em>str</em>) – Path of the model.</p>
-</dd>
-</dl>
-</dd></dl>
-
-<dl class="method">
-<dt id="pytorch_tabnet.abstract_model.TabModel.load_weights_from_unsupervised">
-<code class="sig-name descname">load_weights_from_unsupervised</code><span class="sig-paren">(</span><em class="sig-param">unsupervised_model</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/abstract_model.html#TabModel.load_weights_from_unsupervised"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.load_weights_from_unsupervised" title="Permalink to this definition">¶</a></dt>
-<dd></dd></dl>
-
-<dl class="attribute">
-<dt id="pytorch_tabnet.abstract_model.TabModel.mask_type">
-<code class="sig-name descname">mask_type</code><em class="property">: str</em><em class="property"> = 'sparsemax'</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.mask_type" title="Permalink to this definition">¶</a></dt>
-<dd></dd></dl>
-
-<dl class="attribute">
-<dt id="pytorch_tabnet.abstract_model.TabModel.momentum">
-<code class="sig-name descname">momentum</code><em class="property">: float</em><em class="property"> = 0.02</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.momentum" title="Permalink to this definition">¶</a></dt>
-<dd></dd></dl>
-
-<dl class="attribute">
-<dt id="pytorch_tabnet.abstract_model.TabModel.n_a">
-<code class="sig-name descname">n_a</code><em class="property">: int</em><em class="property"> = 8</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.n_a" title="Permalink to this definition">¶</a></dt>
-<dd></dd></dl>
-
-<dl class="attribute">
-<dt id="pytorch_tabnet.abstract_model.TabModel.n_d">
-<code class="sig-name descname">n_d</code><em class="property">: int</em><em class="property"> = 8</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.n_d" title="Permalink to this definition">¶</a></dt>
-<dd></dd></dl>
-
-<dl class="attribute">
-<dt id="pytorch_tabnet.abstract_model.TabModel.n_independent">
-<code class="sig-name descname">n_independent</code><em class="property">: int</em><em class="property"> = 2</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.n_independent" title="Permalink to this definition">¶</a></dt>
-<dd></dd></dl>
-
-<dl class="attribute">
-<dt id="pytorch_tabnet.abstract_model.TabModel.n_shared">
-<code class="sig-name descname">n_shared</code><em class="property">: int</em><em class="property"> = 2</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.n_shared" title="Permalink to this definition">¶</a></dt>
-<dd></dd></dl>
-
-<dl class="attribute">
-<dt id="pytorch_tabnet.abstract_model.TabModel.n_steps">
-<code class="sig-name descname">n_steps</code><em class="property">: int</em><em class="property"> = 3</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.n_steps" title="Permalink to this definition">¶</a></dt>
-<dd></dd></dl>
-
-<dl class="attribute">
-<dt id="pytorch_tabnet.abstract_model.TabModel.optimizer_fn">
-<code class="sig-name descname">optimizer_fn</code><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.optimizer_fn" title="Permalink to this definition">¶</a></dt>
-<dd><p>alias of <code class="xref py py-class docutils literal notranslate"><span class="pre">torch.optim.adam.Adam</span></code></p>
-</dd></dl>
-
-<dl class="attribute">
-<dt id="pytorch_tabnet.abstract_model.TabModel.optimizer_params">
-<code class="sig-name descname">optimizer_params</code><em class="property">: Dict</em><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.optimizer_params" title="Permalink to this definition">¶</a></dt>
-<dd></dd></dl>
-
-<dl class="attribute">
-<dt id="pytorch_tabnet.abstract_model.TabModel.output_dim">
-<code class="sig-name descname">output_dim</code><em class="property">: int</em><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.output_dim" title="Permalink to this definition">¶</a></dt>
-<dd></dd></dl>
-
-<dl class="method">
-<dt id="pytorch_tabnet.abstract_model.TabModel.predict">
-<code class="sig-name descname">predict</code><span class="sig-paren">(</span><em class="sig-param">X</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/abstract_model.html#TabModel.predict"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.predict" title="Permalink to this definition">¶</a></dt>
-<dd><p>Make predictions on a batch (valid)</p>
-<dl class="field-list simple">
-<dt class="field-odd">Parameters</dt>
-<dd class="field-odd"><p><strong>X</strong> (a :tensor: <cite>torch.Tensor</cite>) – Input data</p>
-</dd>
-<dt class="field-even">Returns</dt>
-<dd class="field-even"><p><strong>predictions</strong> – Predictions of the regression problem</p>
-</dd>
-<dt class="field-odd">Return type</dt>
-<dd class="field-odd"><p>np.array</p>
-</dd>
-</dl>
-</dd></dl>
-
-<dl class="method">
-<dt id="pytorch_tabnet.abstract_model.TabModel.prepare_target">
-<em class="property">abstract </em><code class="sig-name descname">prepare_target</code><span class="sig-paren">(</span><em class="sig-param">y</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/abstract_model.html#TabModel.prepare_target"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.prepare_target" title="Permalink to this definition">¶</a></dt>
-<dd><p>Prepare target before training.</p>
-<dl class="field-list simple">
-<dt class="field-odd">Parameters</dt>
-<dd class="field-odd"><p><strong>y</strong> (a :tensor: <cite>torch.Tensor</cite>) – Target matrix.</p>
-</dd>
-<dt class="field-even">Returns</dt>
-<dd class="field-even"><p>Converted target matrix.</p>
-</dd>
-<dt class="field-odd">Return type</dt>
-<dd class="field-odd"><p><cite>torch.Tensor</cite></p>
-</dd>
-</dl>
-</dd></dl>
-
-<dl class="method">
-<dt id="pytorch_tabnet.abstract_model.TabModel.save_model">
-<code class="sig-name descname">save_model</code><span class="sig-paren">(</span><em class="sig-param">path</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/abstract_model.html#TabModel.save_model"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.save_model" title="Permalink to this definition">¶</a></dt>
-<dd><p>Saving TabNet model in two distinct files.</p>
-<dl class="field-list simple">
-<dt class="field-odd">Parameters</dt>
-<dd class="field-odd"><p><strong>path</strong> (<em>str</em>) – Path of the model.</p>
-</dd>
-<dt class="field-even">Returns</dt>
-<dd class="field-even"><p>input filepath with “.zip” appended</p>
-</dd>
-<dt class="field-odd">Return type</dt>
-<dd class="field-odd"><p>str</p>
-</dd>
-</dl>
-</dd></dl>
-
-<dl class="attribute">
-<dt id="pytorch_tabnet.abstract_model.TabModel.scheduler_fn">
-<code class="sig-name descname">scheduler_fn</code><em class="property">: Any</em><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.scheduler_fn" title="Permalink to this definition">¶</a></dt>
-<dd></dd></dl>
-
-<dl class="attribute">
-<dt id="pytorch_tabnet.abstract_model.TabModel.scheduler_params">
-<code class="sig-name descname">scheduler_params</code><em class="property">: Dict</em><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.scheduler_params" title="Permalink to this definition">¶</a></dt>
-<dd></dd></dl>
-
-<dl class="attribute">
-<dt id="pytorch_tabnet.abstract_model.TabModel.seed">
-<code class="sig-name descname">seed</code><em class="property">: int</em><em class="property"> = 0</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.seed" title="Permalink to this definition">¶</a></dt>
+<dl class="function">
+<dt id="pytorch_tabnet.sparsemax.entmoid15">
+<code class="sig-prename descclassname">pytorch_tabnet.sparsemax.</code><code class="sig-name descname">entmoid15</code><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="headerlink" href="#pytorch_tabnet.sparsemax.entmoid15" title="Permalink to this definition">¶</a></dt>
 <dd></dd></dl>
 
-<dl class="method">
-<dt id="pytorch_tabnet.abstract_model.TabModel.update_fit_params">
-<em class="property">abstract </em><code class="sig-name descname">update_fit_params</code><span class="sig-paren">(</span><em class="sig-param">X_train</em>, <em class="sig-param">y_train</em>, <em class="sig-param">eval_set</em>, <em class="sig-param">weights</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/abstract_model.html#TabModel.update_fit_params"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.update_fit_params" title="Permalink to this definition">¶</a></dt>
-<dd><p>Set attributes relative to fit function.</p>
-<dl class="field-list simple">
-<dt class="field-odd">Parameters</dt>
-<dd class="field-odd"><ul class="simple">
-<li><p><strong>X_train</strong> (<em>np.ndarray</em>) – Train set</p></li>
-<li><p><strong>y_train</strong> (<em>np.array</em>) – Train targets</p></li>
-<li><p><strong>eval_set</strong> (<em>list of tuple</em>) – List of eval tuple set (X, y).</p></li>
-<li><p><strong>weights</strong> (<em>bool</em><em> or </em><em>dictionnary</em>) – 0 for no balancing
-1 for automated balancing</p></li>
-</ul>
-</dd>
-</dl>
-</dd></dl>
-
-<dl class="attribute">
-<dt id="pytorch_tabnet.abstract_model.TabModel.verbose">
-<code class="sig-name descname">verbose</code><em class="property">: int</em><em class="property"> = 1</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.verbose" title="Permalink to this definition">¶</a></dt>
+<dl class="function">
+<dt id="pytorch_tabnet.sparsemax.sparsemax">
+<code class="sig-prename descclassname">pytorch_tabnet.sparsemax.</code><code class="sig-name descname">sparsemax</code><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="headerlink" href="#pytorch_tabnet.sparsemax.sparsemax" title="Permalink to this definition">¶</a></dt>
 <dd></dd></dl>
 
-</dd></dl>
-
-</div>
-<div class="section" id="module-pytorch_tabnet.callbacks">
+</section>
+<section id="module-pytorch_tabnet.callbacks">
 <span id="pytorch-tabnet-callbacks-module"></span><h2>pytorch_tabnet.callbacks module<a class="headerlink" href="#module-pytorch_tabnet.callbacks" title="Permalink to this headline">¶</a></h2>
 <dl class="class">
 <dt id="pytorch_tabnet.callbacks.Callback">
@@ -1812,31 +1388,42 @@ <h3>Multi-class / multi-label utility function<a class="headerlink" href="#multi
 
 </dd></dl>
 
-</div>
-<div class="section" id="module-pytorch_tabnet.pretraining">
-<span id="pytorch-tabnet-pretraining-module"></span><h2>pytorch_tabnet.pretraining module<a class="headerlink" href="#module-pytorch_tabnet.pretraining" title="Permalink to this headline">¶</a></h2>
+</section>
+<section id="module-pytorch_tabnet.abstract_model">
+<span id="pytorch-tabnet-abstract-model-module"></span><h2>pytorch_tabnet.abstract_model module<a class="headerlink" href="#module-pytorch_tabnet.abstract_model" title="Permalink to this headline">¶</a></h2>
 <dl class="class">
-<dt id="pytorch_tabnet.pretraining.TabNetPretrainer">
-<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.pretraining.</code><code class="sig-name descname">TabNetPretrainer</code><span class="sig-paren">(</span><em class="sig-param">n_d: int = 8</em>, <em class="sig-param">n_a: int = 8</em>, <em class="sig-param">n_steps: int = 3</em>, <em class="sig-param">gamma: float = 1.3</em>, <em class="sig-param">cat_idxs: List[int] = &lt;factory&gt;</em>, <em class="sig-param">cat_dims: List[int] = &lt;factory&gt;</em>, <em class="sig-param">cat_emb_dim: int = 1</em>, <em class="sig-param">n_independent: int = 2</em>, <em class="sig-param">n_shared: int = 2</em>, <em class="sig-param">epsilon: float = 1e-15</em>, <em class="sig-param">momentum: float = 0.02</em>, <em class="sig-param">lambda_sparse: float = 0.001</em>, <em class="sig-param">seed: int = 0</em>, <em class="sig-param">clip_value: int = 1</em>, <em class="sig-param">verbose: int = 1</em>, <em class="sig-param">optimizer_fn: Any = &lt;class 'torch.optim.adam.Adam'&gt;</em>, <em class="sig-param">optimizer_params: Dict = &lt;factory&gt;</em>, <em class="sig-param">scheduler_fn: Any = None</em>, <em class="sig-param">scheduler_params: Dict = &lt;factory&gt;</em>, <em class="sig-param">mask_type: str = 'sparsemax'</em>, <em class="sig-param">input_dim: int = None</em>, <em class="sig-param">output_dim: int = None</em>, <em class="sig-param">device_name: str = 'auto'</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/pretraining.html#TabNetPretrainer"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.pretraining.TabNetPretrainer" title="Permalink to this definition">¶</a></dt>
-<dd><p>Bases: <a class="reference internal" href="#pytorch_tabnet.abstract_model.TabModel" title="pytorch_tabnet.abstract_model.TabModel"><code class="xref py py-class docutils literal notranslate"><span class="pre">pytorch_tabnet.abstract_model.TabModel</span></code></a></p>
+<dt id="pytorch_tabnet.abstract_model.TabModel">
+<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.abstract_model.</code><code class="sig-name descname">TabModel</code><span class="sig-paren">(</span><em class="sig-param">n_d: int = 8</em>, <em class="sig-param">n_a: int = 8</em>, <em class="sig-param">n_steps: int = 3</em>, <em class="sig-param">gamma: float = 1.3</em>, <em class="sig-param">cat_idxs: List[int] = &lt;factory&gt;</em>, <em class="sig-param">cat_dims: List[int] = &lt;factory&gt;</em>, <em class="sig-param">cat_emb_dim: int = 1</em>, <em class="sig-param">n_independent: int = 2</em>, <em class="sig-param">n_shared: int = 2</em>, <em class="sig-param">epsilon: float = 1e-15</em>, <em class="sig-param">momentum: float = 0.02</em>, <em class="sig-param">lambda_sparse: float = 0.001</em>, <em class="sig-param">seed: int = 0</em>, <em class="sig-param">clip_value: int = 1</em>, <em class="sig-param">verbose: int = 1</em>, <em class="sig-param">optimizer_fn: Any = &lt;class 'torch.optim.adam.Adam'&gt;</em>, <em class="sig-param">optimizer_params: Dict = &lt;factory&gt;</em>, <em class="sig-param">scheduler_fn: Any = None</em>, <em class="sig-param">scheduler_params: Dict = &lt;factory&gt;</em>, <em class="sig-param">mask_type: str = 'sparsemax'</em>, <em class="sig-param">input_dim: int = None</em>, <em class="sig-param">output_dim: int = None</em>, <em class="sig-param">device_name: str = 'auto'</em>, <em class="sig-param">n_shared_decoder: int = 1</em>, <em class="sig-param">n_indep_decoder: int = 1</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/abstract_model.html#TabModel"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel" title="Permalink to this definition">¶</a></dt>
+<dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">sklearn.base.BaseEstimator</span></code></p>
+<p>Class for TabNet model.</p>
 <dl class="attribute">
-<dt id="pytorch_tabnet.pretraining.TabNetPretrainer.cat_dims">
-<code class="sig-name descname">cat_dims</code><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.pretraining.TabNetPretrainer.cat_dims" title="Permalink to this definition">¶</a></dt>
+<dt id="pytorch_tabnet.abstract_model.TabModel.cat_dims">
+<code class="sig-name descname">cat_dims</code><em class="property">: List[int]</em><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.cat_dims" title="Permalink to this definition">¶</a></dt>
 <dd></dd></dl>
 
 <dl class="attribute">
-<dt id="pytorch_tabnet.pretraining.TabNetPretrainer.cat_idxs">
-<code class="sig-name descname">cat_idxs</code><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.pretraining.TabNetPretrainer.cat_idxs" title="Permalink to this definition">¶</a></dt>
+<dt id="pytorch_tabnet.abstract_model.TabModel.cat_emb_dim">
+<code class="sig-name descname">cat_emb_dim</code><em class="property">: int</em><em class="property"> = 1</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.cat_emb_dim" title="Permalink to this definition">¶</a></dt>
 <dd></dd></dl>
 
-<dl class="method">
-<dt id="pytorch_tabnet.pretraining.TabNetPretrainer.compute_loss">
-<code class="sig-name descname">compute_loss</code><span class="sig-paren">(</span><em class="sig-param">output</em>, <em class="sig-param">embedded_x</em>, <em class="sig-param">obf_vars</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/pretraining.html#TabNetPretrainer.compute_loss"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.pretraining.TabNetPretrainer.compute_loss" title="Permalink to this definition">¶</a></dt>
-<dd><p>Compute the loss.</p>
-<dl class="field-list simple">
-<dt class="field-odd">Parameters</dt>
-<dd class="field-odd"><ul class="simple">
-<li><p><strong>y_score</strong> (a :tensor: <cite>torch.Tensor</cite>) – Score matrix</p></li>
+<dl class="attribute">
+<dt id="pytorch_tabnet.abstract_model.TabModel.cat_idxs">
+<code class="sig-name descname">cat_idxs</code><em class="property">: List[int]</em><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.cat_idxs" title="Permalink to this definition">¶</a></dt>
+<dd></dd></dl>
+
+<dl class="attribute">
+<dt id="pytorch_tabnet.abstract_model.TabModel.clip_value">
+<code class="sig-name descname">clip_value</code><em class="property">: int</em><em class="property"> = 1</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.clip_value" title="Permalink to this definition">¶</a></dt>
+<dd></dd></dl>
+
+<dl class="method">
+<dt id="pytorch_tabnet.abstract_model.TabModel.compute_loss">
+<em class="property">abstract </em><code class="sig-name descname">compute_loss</code><span class="sig-paren">(</span><em class="sig-param">y_score</em>, <em class="sig-param">y_true</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/abstract_model.html#TabModel.compute_loss"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.compute_loss" title="Permalink to this definition">¶</a></dt>
+<dd><p>Compute the loss.</p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters</dt>
+<dd class="field-odd"><ul class="simple">
+<li><p><strong>y_score</strong> (a :tensor: <cite>torch.Tensor</cite>) – Score matrix</p></li>
 <li><p><strong>y_true</strong> (a :tensor: <cite>torch.Tensor</cite>) – Target matrix</p></li>
 </ul>
 </dd>
@@ -1849,25 +1436,57 @@ <h3>Multi-class / multi-label utility function<a class="headerlink" href="#multi
 </dl>
 </dd></dl>
 
+<dl class="attribute">
+<dt id="pytorch_tabnet.abstract_model.TabModel.device_name">
+<code class="sig-name descname">device_name</code><em class="property">: str</em><em class="property"> = 'auto'</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.device_name" title="Permalink to this definition">¶</a></dt>
+<dd></dd></dl>
+
+<dl class="attribute">
+<dt id="pytorch_tabnet.abstract_model.TabModel.epsilon">
+<code class="sig-name descname">epsilon</code><em class="property">: float</em><em class="property"> = 1e-15</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.epsilon" title="Permalink to this definition">¶</a></dt>
+<dd></dd></dl>
+
 <dl class="method">
-<dt id="pytorch_tabnet.pretraining.TabNetPretrainer.fit">
-<code class="sig-name descname">fit</code><span class="sig-paren">(</span><em class="sig-param">X_train</em>, <em class="sig-param">eval_set=None</em>, <em class="sig-param">eval_name=None</em>, <em class="sig-param">loss_fn=None</em>, <em class="sig-param">pretraining_ratio=0.5</em>, <em class="sig-param">weights=0</em>, <em class="sig-param">max_epochs=100</em>, <em class="sig-param">patience=10</em>, <em class="sig-param">batch_size=1024</em>, <em class="sig-param">virtual_batch_size=128</em>, <em class="sig-param">num_workers=0</em>, <em class="sig-param">drop_last=False</em>, <em class="sig-param">callbacks=None</em>, <em class="sig-param">pin_memory=True</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/pretraining.html#TabNetPretrainer.fit"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.pretraining.TabNetPretrainer.fit" title="Permalink to this definition">¶</a></dt>
+<dt id="pytorch_tabnet.abstract_model.TabModel.explain">
+<code class="sig-name descname">explain</code><span class="sig-paren">(</span><em class="sig-param">X</em>, <em class="sig-param">normalize=False</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/abstract_model.html#TabModel.explain"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.explain" title="Permalink to this definition">¶</a></dt>
+<dd><p>Return local explanation</p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters</dt>
+<dd class="field-odd"><ul class="simple">
+<li><p><strong>X</strong> (tensor: <cite>torch.Tensor</cite>) – Input data</p></li>
+<li><p><strong>normalize</strong> (<em>bool</em><em> (</em><em>default False</em><em>)</em>) – Wheter to normalize so that sum of features are equal to 1</p></li>
+</ul>
+</dd>
+<dt class="field-even">Returns</dt>
+<dd class="field-even"><p><ul class="simple">
+<li><p><strong>M_explain</strong> (<em>matrix</em>) – Importance per sample, per columns.</p></li>
+<li><p><strong>masks</strong> (<em>matrix</em>) – Sparse matrix showing attention masks used by network.</p></li>
+</ul>
+</p>
+</dd>
+</dl>
+</dd></dl>
+
+<dl class="method">
+<dt id="pytorch_tabnet.abstract_model.TabModel.fit">
+<code class="sig-name descname">fit</code><span class="sig-paren">(</span><em class="sig-param">X_train</em>, <em class="sig-param">y_train</em>, <em class="sig-param">eval_set=None</em>, <em class="sig-param">eval_name=None</em>, <em class="sig-param">eval_metric=None</em>, <em class="sig-param">loss_fn=None</em>, <em class="sig-param">weights=0</em>, <em class="sig-param">max_epochs=100</em>, <em class="sig-param">patience=10</em>, <em class="sig-param">batch_size=1024</em>, <em class="sig-param">virtual_batch_size=128</em>, <em class="sig-param">num_workers=0</em>, <em class="sig-param">drop_last=True</em>, <em class="sig-param">callbacks=None</em>, <em class="sig-param">pin_memory=True</em>, <em class="sig-param">from_unsupervised=None</em>, <em class="sig-param">warm_start=False</em>, <em class="sig-param">augmentations=None</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/abstract_model.html#TabModel.fit"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.fit" title="Permalink to this definition">¶</a></dt>
 <dd><p>Train a neural network stored in self.network
 Using train_dataloader for training data and
 valid_dataloader for validation.</p>
 <dl class="field-list simple">
 <dt class="field-odd">Parameters</dt>
 <dd class="field-odd"><ul class="simple">
-<li><p><strong>X_train</strong> (<em>np.ndarray</em>) – Train set to reconstruct in self supervision</p></li>
-<li><p><strong>eval_set</strong> (<em>list of np.array</em>) – List of evaluation set
+<li><p><strong>X_train</strong> (<em>np.ndarray</em>) – Train set</p></li>
+<li><p><strong>y_train</strong> (<em>np.array</em>) – Train targets</p></li>
+<li><p><strong>eval_set</strong> (<em>list of tuple</em>) – List of eval tuple set (X, y).
 The last one is used for early stopping</p></li>
 <li><p><strong>eval_name</strong> (<em>list of str</em>) – List of eval set names.</p></li>
 <li><p><strong>eval_metric</strong> (<em>list of str</em>) – List of evaluation metrics.
 The last metric is used for early stopping.</p></li>
-<li><p><strong>loss_fn</strong> (<em>callable</em><em> or </em><em>None</em>) – a PyTorch loss function
-should be left to None for self supervised and non experts</p></li>
-<li><p><strong>pretraining_ratio</strong> (<em>float</em>) – Between 0 and 1, percentage of feature to mask for reconstruction</p></li>
-<li><p><strong>weights</strong> (<em>np.array</em>) – Sampling weights for each example.</p></li>
+<li><p><strong>loss_fn</strong> (<em>callable</em><em> or </em><em>None</em>) – a PyTorch loss function</p></li>
+<li><p><strong>weights</strong> (<em>bool</em><em> or </em><em>dictionnary</em>) – 0 for no balancing
+1 for automated balancing
+dict for custom weights per class</p></li>
 <li><p><strong>max_epochs</strong> (<em>int</em>) – Maximum number of epochs during training</p></li>
 <li><p><strong>patience</strong> (<em>int</em>) – Number of consecutive non improving epoch before early stopping</p></li>
 <li><p><strong>batch_size</strong> (<em>int</em>) – Training batch size</p></li>
@@ -1876,19 +1495,113 @@ <h3>Multi-class / multi-label utility function<a class="headerlink" href="#multi
 <li><p><strong>drop_last</strong> (<em>bool</em>) – Whether to drop last batch during training</p></li>
 <li><p><strong>callbacks</strong> (<em>list of callback function</em>) – List of custom callbacks</p></li>
 <li><p><strong>pin_memory</strong> (<em>bool</em>) – Whether to set pin_memory to True or False during training</p></li>
+<li><p><strong>from_unsupervised</strong> (<em>unsupervised trained model</em>) – Use a previously self supervised model as starting weights</p></li>
+<li><p><strong>warm_start</strong> (<em>bool</em>) – If True, current model parameters are used to start training</p></li>
 </ul>
 </dd>
 </dl>
 </dd></dl>
 
 <dl class="attribute">
-<dt id="pytorch_tabnet.pretraining.TabNetPretrainer.optimizer_params">
-<code class="sig-name descname">optimizer_params</code><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.pretraining.TabNetPretrainer.optimizer_params" title="Permalink to this definition">¶</a></dt>
+<dt id="pytorch_tabnet.abstract_model.TabModel.gamma">
+<code class="sig-name descname">gamma</code><em class="property">: float</em><em class="property"> = 1.3</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.gamma" title="Permalink to this definition">¶</a></dt>
+<dd></dd></dl>
+
+<dl class="attribute">
+<dt id="pytorch_tabnet.abstract_model.TabModel.input_dim">
+<code class="sig-name descname">input_dim</code><em class="property">: int</em><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.input_dim" title="Permalink to this definition">¶</a></dt>
+<dd></dd></dl>
+
+<dl class="attribute">
+<dt id="pytorch_tabnet.abstract_model.TabModel.lambda_sparse">
+<code class="sig-name descname">lambda_sparse</code><em class="property">: float</em><em class="property"> = 0.001</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.lambda_sparse" title="Permalink to this definition">¶</a></dt>
 <dd></dd></dl>
 
 <dl class="method">
-<dt id="pytorch_tabnet.pretraining.TabNetPretrainer.predict">
-<code class="sig-name descname">predict</code><span class="sig-paren">(</span><em class="sig-param">X</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/pretraining.html#TabNetPretrainer.predict"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.pretraining.TabNetPretrainer.predict" title="Permalink to this definition">¶</a></dt>
+<dt id="pytorch_tabnet.abstract_model.TabModel.load_class_attrs">
+<code class="sig-name descname">load_class_attrs</code><span class="sig-paren">(</span><em class="sig-param">class_attrs</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/abstract_model.html#TabModel.load_class_attrs"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.load_class_attrs" title="Permalink to this definition">¶</a></dt>
+<dd></dd></dl>
+
+<dl class="method">
+<dt id="pytorch_tabnet.abstract_model.TabModel.load_model">
+<code class="sig-name descname">load_model</code><span class="sig-paren">(</span><em class="sig-param">filepath</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/abstract_model.html#TabModel.load_model"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.load_model" title="Permalink to this definition">¶</a></dt>
+<dd><p>Load TabNet model.</p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters</dt>
+<dd class="field-odd"><p><strong>filepath</strong> (<em>str</em>) – Path of the model.</p>
+</dd>
+</dl>
+</dd></dl>
+
+<dl class="method">
+<dt id="pytorch_tabnet.abstract_model.TabModel.load_weights_from_unsupervised">
+<code class="sig-name descname">load_weights_from_unsupervised</code><span class="sig-paren">(</span><em class="sig-param">unsupervised_model</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/abstract_model.html#TabModel.load_weights_from_unsupervised"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.load_weights_from_unsupervised" title="Permalink to this definition">¶</a></dt>
+<dd></dd></dl>
+
+<dl class="attribute">
+<dt id="pytorch_tabnet.abstract_model.TabModel.mask_type">
+<code class="sig-name descname">mask_type</code><em class="property">: str</em><em class="property"> = 'sparsemax'</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.mask_type" title="Permalink to this definition">¶</a></dt>
+<dd></dd></dl>
+
+<dl class="attribute">
+<dt id="pytorch_tabnet.abstract_model.TabModel.momentum">
+<code class="sig-name descname">momentum</code><em class="property">: float</em><em class="property"> = 0.02</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.momentum" title="Permalink to this definition">¶</a></dt>
+<dd></dd></dl>
+
+<dl class="attribute">
+<dt id="pytorch_tabnet.abstract_model.TabModel.n_a">
+<code class="sig-name descname">n_a</code><em class="property">: int</em><em class="property"> = 8</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.n_a" title="Permalink to this definition">¶</a></dt>
+<dd></dd></dl>
+
+<dl class="attribute">
+<dt id="pytorch_tabnet.abstract_model.TabModel.n_d">
+<code class="sig-name descname">n_d</code><em class="property">: int</em><em class="property"> = 8</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.n_d" title="Permalink to this definition">¶</a></dt>
+<dd></dd></dl>
+
+<dl class="attribute">
+<dt id="pytorch_tabnet.abstract_model.TabModel.n_indep_decoder">
+<code class="sig-name descname">n_indep_decoder</code><em class="property">: int</em><em class="property"> = 1</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.n_indep_decoder" title="Permalink to this definition">¶</a></dt>
+<dd></dd></dl>
+
+<dl class="attribute">
+<dt id="pytorch_tabnet.abstract_model.TabModel.n_independent">
+<code class="sig-name descname">n_independent</code><em class="property">: int</em><em class="property"> = 2</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.n_independent" title="Permalink to this definition">¶</a></dt>
+<dd></dd></dl>
+
+<dl class="attribute">
+<dt id="pytorch_tabnet.abstract_model.TabModel.n_shared">
+<code class="sig-name descname">n_shared</code><em class="property">: int</em><em class="property"> = 2</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.n_shared" title="Permalink to this definition">¶</a></dt>
+<dd></dd></dl>
+
+<dl class="attribute">
+<dt id="pytorch_tabnet.abstract_model.TabModel.n_shared_decoder">
+<code class="sig-name descname">n_shared_decoder</code><em class="property">: int</em><em class="property"> = 1</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.n_shared_decoder" title="Permalink to this definition">¶</a></dt>
+<dd></dd></dl>
+
+<dl class="attribute">
+<dt id="pytorch_tabnet.abstract_model.TabModel.n_steps">
+<code class="sig-name descname">n_steps</code><em class="property">: int</em><em class="property"> = 3</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.n_steps" title="Permalink to this definition">¶</a></dt>
+<dd></dd></dl>
+
+<dl class="attribute">
+<dt id="pytorch_tabnet.abstract_model.TabModel.optimizer_fn">
+<code class="sig-name descname">optimizer_fn</code><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.optimizer_fn" title="Permalink to this definition">¶</a></dt>
+<dd><p>alias of <code class="xref py py-class docutils literal notranslate"><span class="pre">torch.optim.adam.Adam</span></code></p>
+</dd></dl>
+
+<dl class="attribute">
+<dt id="pytorch_tabnet.abstract_model.TabModel.optimizer_params">
+<code class="sig-name descname">optimizer_params</code><em class="property">: Dict</em><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.optimizer_params" title="Permalink to this definition">¶</a></dt>
+<dd></dd></dl>
+
+<dl class="attribute">
+<dt id="pytorch_tabnet.abstract_model.TabModel.output_dim">
+<code class="sig-name descname">output_dim</code><em class="property">: int</em><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.output_dim" title="Permalink to this definition">¶</a></dt>
+<dd></dd></dl>
+
+<dl class="method">
+<dt id="pytorch_tabnet.abstract_model.TabModel.predict">
+<code class="sig-name descname">predict</code><span class="sig-paren">(</span><em class="sig-param">X</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/abstract_model.html#TabModel.predict"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.predict" title="Permalink to this definition">¶</a></dt>
 <dd><p>Make predictions on a batch (valid)</p>
 <dl class="field-list simple">
 <dt class="field-odd">Parameters</dt>
@@ -1904,8 +1617,8 @@ <h3>Multi-class / multi-label utility function<a class="headerlink" href="#multi
 </dd></dl>
 
 <dl class="method">
-<dt id="pytorch_tabnet.pretraining.TabNetPretrainer.prepare_target">
-<code class="sig-name descname">prepare_target</code><span class="sig-paren">(</span><em class="sig-param">y</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/pretraining.html#TabNetPretrainer.prepare_target"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.pretraining.TabNetPretrainer.prepare_target" title="Permalink to this definition">¶</a></dt>
+<dt id="pytorch_tabnet.abstract_model.TabModel.prepare_target">
+<em class="property">abstract </em><code class="sig-name descname">prepare_target</code><span class="sig-paren">(</span><em class="sig-param">y</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/abstract_model.html#TabModel.prepare_target"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.prepare_target" title="Permalink to this definition">¶</a></dt>
 <dd><p>Prepare target before training.</p>
 <dl class="field-list simple">
 <dt class="field-odd">Parameters</dt>
@@ -1920,19 +1633,41 @@ <h3>Multi-class / multi-label utility function<a class="headerlink" href="#multi
 </dl>
 </dd></dl>
 
+<dl class="method">
+<dt id="pytorch_tabnet.abstract_model.TabModel.save_model">
+<code class="sig-name descname">save_model</code><span class="sig-paren">(</span><em class="sig-param">path</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/abstract_model.html#TabModel.save_model"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.save_model" title="Permalink to this definition">¶</a></dt>
+<dd><p>Saving TabNet model in two distinct files.</p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters</dt>
+<dd class="field-odd"><p><strong>path</strong> (<em>str</em>) – Path of the model.</p>
+</dd>
+<dt class="field-even">Returns</dt>
+<dd class="field-even"><p>input filepath with “.zip” appended</p>
+</dd>
+<dt class="field-odd">Return type</dt>
+<dd class="field-odd"><p>str</p>
+</dd>
+</dl>
+</dd></dl>
+
 <dl class="attribute">
-<dt id="pytorch_tabnet.pretraining.TabNetPretrainer.scheduler_params">
-<code class="sig-name descname">scheduler_params</code><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.pretraining.TabNetPretrainer.scheduler_params" title="Permalink to this definition">¶</a></dt>
+<dt id="pytorch_tabnet.abstract_model.TabModel.scheduler_fn">
+<code class="sig-name descname">scheduler_fn</code><em class="property">: Any</em><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.scheduler_fn" title="Permalink to this definition">¶</a></dt>
 <dd></dd></dl>
 
-<dl class="method">
-<dt id="pytorch_tabnet.pretraining.TabNetPretrainer.stack_batches">
-<code class="sig-name descname">stack_batches</code><span class="sig-paren">(</span><em class="sig-param">list_output</em>, <em class="sig-param">list_embedded_x</em>, <em class="sig-param">list_obfuscation</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/pretraining.html#TabNetPretrainer.stack_batches"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.pretraining.TabNetPretrainer.stack_batches" title="Permalink to this definition">¶</a></dt>
+<dl class="attribute">
+<dt id="pytorch_tabnet.abstract_model.TabModel.scheduler_params">
+<code class="sig-name descname">scheduler_params</code><em class="property">: Dict</em><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.scheduler_params" title="Permalink to this definition">¶</a></dt>
+<dd></dd></dl>
+
+<dl class="attribute">
+<dt id="pytorch_tabnet.abstract_model.TabModel.seed">
+<code class="sig-name descname">seed</code><em class="property">: int</em><em class="property"> = 0</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.seed" title="Permalink to this definition">¶</a></dt>
 <dd></dd></dl>
 
 <dl class="method">
-<dt id="pytorch_tabnet.pretraining.TabNetPretrainer.update_fit_params">
-<code class="sig-name descname">update_fit_params</code><span class="sig-paren">(</span><em class="sig-param">weights</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/pretraining.html#TabNetPretrainer.update_fit_params"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.pretraining.TabNetPretrainer.update_fit_params" title="Permalink to this definition">¶</a></dt>
+<dt id="pytorch_tabnet.abstract_model.TabModel.update_fit_params">
+<em class="property">abstract </em><code class="sig-name descname">update_fit_params</code><span class="sig-paren">(</span><em class="sig-param">X_train</em>, <em class="sig-param">y_train</em>, <em class="sig-param">eval_set</em>, <em class="sig-param">weights</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/abstract_model.html#TabModel.update_fit_params"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.update_fit_params" title="Permalink to this definition">¶</a></dt>
 <dd><p>Set attributes relative to fit function.</p>
 <dl class="field-list simple">
 <dt class="field-odd">Parameters</dt>
@@ -1947,164 +1682,392 @@ <h3>Multi-class / multi-label utility function<a class="headerlink" href="#multi
 </dl>
 </dd></dl>
 
+<dl class="attribute">
+<dt id="pytorch_tabnet.abstract_model.TabModel.verbose">
+<code class="sig-name descname">verbose</code><em class="property">: int</em><em class="property"> = 1</em><a class="headerlink" href="#pytorch_tabnet.abstract_model.TabModel.verbose" title="Permalink to this definition">¶</a></dt>
+<dd></dd></dl>
+
 </dd></dl>
 
-</div>
-<div class="section" id="module-pytorch_tabnet.sparsemax">
-<span id="pytorch-tabnet-sparsemax-module"></span><h2>pytorch_tabnet.sparsemax module<a class="headerlink" href="#module-pytorch_tabnet.sparsemax" title="Permalink to this headline">¶</a></h2>
+</section>
+<section id="module-pytorch_tabnet.pretraining">
+<span id="pytorch-tabnet-pretraining-module"></span><h2>pytorch_tabnet.pretraining module<a class="headerlink" href="#module-pytorch_tabnet.pretraining" title="Permalink to this headline">¶</a></h2>
 <dl class="class">
-<dt id="pytorch_tabnet.sparsemax.Entmax15">
-<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.sparsemax.</code><code class="sig-name descname">Entmax15</code><span class="sig-paren">(</span><em class="sig-param">dim=-1</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/sparsemax.html#Entmax15"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.sparsemax.Entmax15" title="Permalink to this definition">¶</a></dt>
-<dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">torch.nn.modules.module.Module</span></code></p>
-<dl class="method">
-<dt id="pytorch_tabnet.sparsemax.Entmax15.forward">
-<code class="sig-name descname">forward</code><span class="sig-paren">(</span><em class="sig-param">input</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/sparsemax.html#Entmax15.forward"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.sparsemax.Entmax15.forward" title="Permalink to this definition">¶</a></dt>
-<dd><p>Defines the computation performed at every call.</p>
-<p>Should be overridden by all subclasses.</p>
-<div class="admonition note">
-<p class="admonition-title">Note</p>
-<p>Although the recipe for forward pass needs to be defined within
-this function, one should call the <code class="xref py py-class docutils literal notranslate"><span class="pre">Module</span></code> instance afterwards
-instead of this since the former takes care of running the
-registered hooks while the latter silently ignores them.</p>
-</div>
-</dd></dl>
+<dt id="pytorch_tabnet.pretraining.TabNetPretrainer">
+<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.pretraining.</code><code class="sig-name descname">TabNetPretrainer</code><span class="sig-paren">(</span><em class="sig-param">n_d: int = 8</em>, <em class="sig-param">n_a: int = 8</em>, <em class="sig-param">n_steps: int = 3</em>, <em class="sig-param">gamma: float = 1.3</em>, <em class="sig-param">cat_idxs: List[int] = &lt;factory&gt;</em>, <em class="sig-param">cat_dims: List[int] = &lt;factory&gt;</em>, <em class="sig-param">cat_emb_dim: int = 1</em>, <em class="sig-param">n_independent: int = 2</em>, <em class="sig-param">n_shared: int = 2</em>, <em class="sig-param">epsilon: float = 1e-15</em>, <em class="sig-param">momentum: float = 0.02</em>, <em class="sig-param">lambda_sparse: float = 0.001</em>, <em class="sig-param">seed: int = 0</em>, <em class="sig-param">clip_value: int = 1</em>, <em class="sig-param">verbose: int = 1</em>, <em class="sig-param">optimizer_fn: Any = &lt;class 'torch.optim.adam.Adam'&gt;</em>, <em class="sig-param">optimizer_params: Dict = &lt;factory&gt;</em>, <em class="sig-param">scheduler_fn: Any = None</em>, <em class="sig-param">scheduler_params: Dict = &lt;factory&gt;</em>, <em class="sig-param">mask_type: str = 'sparsemax'</em>, <em class="sig-param">input_dim: int = None</em>, <em class="sig-param">output_dim: int = None</em>, <em class="sig-param">device_name: str = 'auto'</em>, <em class="sig-param">n_shared_decoder: int = 1</em>, <em class="sig-param">n_indep_decoder: int = 1</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/pretraining.html#TabNetPretrainer"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.pretraining.TabNetPretrainer" title="Permalink to this definition">¶</a></dt>
+<dd><p>Bases: <a class="reference internal" href="#pytorch_tabnet.abstract_model.TabModel" title="pytorch_tabnet.abstract_model.TabModel"><code class="xref py py-class docutils literal notranslate"><span class="pre">pytorch_tabnet.abstract_model.TabModel</span></code></a></p>
+<dl class="attribute">
+<dt id="pytorch_tabnet.pretraining.TabNetPretrainer.cat_dims">
+<code class="sig-name descname">cat_dims</code><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.pretraining.TabNetPretrainer.cat_dims" title="Permalink to this definition">¶</a></dt>
+<dd></dd></dl>
 
-</dd></dl>
+<dl class="attribute">
+<dt id="pytorch_tabnet.pretraining.TabNetPretrainer.cat_idxs">
+<code class="sig-name descname">cat_idxs</code><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.pretraining.TabNetPretrainer.cat_idxs" title="Permalink to this definition">¶</a></dt>
+<dd></dd></dl>
 
-<dl class="class">
-<dt id="pytorch_tabnet.sparsemax.Entmax15Function">
-<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.sparsemax.</code><code class="sig-name descname">Entmax15Function</code><a class="reference internal" href="../_modules/pytorch_tabnet/sparsemax.html#Entmax15Function"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.sparsemax.Entmax15Function" title="Permalink to this definition">¶</a></dt>
-<dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">torch.autograd.function.Function</span></code></p>
-<p>An implementation of exact Entmax with alpha=1.5 (B. Peters, V. Niculae, A. Martins). See
-:cite:<a href="#id1"><span class="problematic" id="id2">`</span></a><a class="reference external" href="https://arxiv.org/abs/1905.05702">https://arxiv.org/abs/1905.05702</a> for detailed description.
-Source: <a class="reference external" href="https://github.com/deep-spin/entmax">https://github.com/deep-spin/entmax</a></p>
 <dl class="method">
-<dt id="pytorch_tabnet.sparsemax.Entmax15Function.backward">
-<em class="property">static </em><code class="sig-name descname">backward</code><span class="sig-paren">(</span><em class="sig-param">ctx</em>, <em class="sig-param">grad_output</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/sparsemax.html#Entmax15Function.backward"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.sparsemax.Entmax15Function.backward" title="Permalink to this definition">¶</a></dt>
-<dd><p>Defines a formula for differentiating the operation.</p>
-<p>This function is to be overridden by all subclasses.</p>
-<p>It must accept a context <code class="xref py py-attr docutils literal notranslate"><span class="pre">ctx</span></code> as the first argument, followed by
-as many outputs did <a class="reference internal" href="#pytorch_tabnet.sparsemax.Entmax15Function.forward" title="pytorch_tabnet.sparsemax.Entmax15Function.forward"><code class="xref py py-func docutils literal notranslate"><span class="pre">forward()</span></code></a> return, and it should return as many
-tensors, as there were inputs to <a class="reference internal" href="#pytorch_tabnet.sparsemax.Entmax15Function.forward" title="pytorch_tabnet.sparsemax.Entmax15Function.forward"><code class="xref py py-func docutils literal notranslate"><span class="pre">forward()</span></code></a>. Each argument is the
-gradient w.r.t the given output, and each returned value should be the
-gradient w.r.t. the corresponding input.</p>
-<p>The context can be used to retrieve tensors saved during the forward
-pass. It also has an attribute <code class="xref py py-attr docutils literal notranslate"><span class="pre">ctx.needs_input_grad</span></code> as a tuple
-of booleans representing whether each input needs gradient. E.g.,
-<a class="reference internal" href="#pytorch_tabnet.sparsemax.Entmax15Function.backward" title="pytorch_tabnet.sparsemax.Entmax15Function.backward"><code class="xref py py-func docutils literal notranslate"><span class="pre">backward()</span></code></a> will have <code class="docutils literal notranslate"><span class="pre">ctx.needs_input_grad[0]</span> <span class="pre">=</span> <span class="pre">True</span></code> if the
-first input to <a class="reference internal" href="#pytorch_tabnet.sparsemax.Entmax15Function.forward" title="pytorch_tabnet.sparsemax.Entmax15Function.forward"><code class="xref py py-func docutils literal notranslate"><span class="pre">forward()</span></code></a> needs gradient computated w.r.t. the
-output.</p>
-</dd></dl>
+<dt id="pytorch_tabnet.pretraining.TabNetPretrainer.compute_loss">
+<code class="sig-name descname">compute_loss</code><span class="sig-paren">(</span><em class="sig-param">output</em>, <em class="sig-param">embedded_x</em>, <em class="sig-param">obf_vars</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/pretraining.html#TabNetPretrainer.compute_loss"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.pretraining.TabNetPretrainer.compute_loss" title="Permalink to this definition">¶</a></dt>
+<dd><p>Compute the loss.</p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters</dt>
+<dd class="field-odd"><ul class="simple">
+<li><p><strong>y_score</strong> (a :tensor: <cite>torch.Tensor</cite>) – Score matrix</p></li>
+<li><p><strong>y_true</strong> (a :tensor: <cite>torch.Tensor</cite>) – Target matrix</p></li>
+</ul>
+</dd>
+<dt class="field-even">Returns</dt>
+<dd class="field-even"><p>Loss value</p>
+</dd>
+<dt class="field-odd">Return type</dt>
+<dd class="field-odd"><p>float</p>
+</dd>
+</dl>
+</dd></dl>
 
 <dl class="method">
-<dt id="pytorch_tabnet.sparsemax.Entmax15Function.forward">
-<em class="property">static </em><code class="sig-name descname">forward</code><span class="sig-paren">(</span><em class="sig-param">ctx</em>, <em class="sig-param">input</em>, <em class="sig-param">dim=-1</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/sparsemax.html#Entmax15Function.forward"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.sparsemax.Entmax15Function.forward" title="Permalink to this definition">¶</a></dt>
-<dd><p>Performs the operation.</p>
-<p>This function is to be overridden by all subclasses.</p>
-<p>It must accept a context ctx as the first argument, followed by any
-number of arguments (tensors or other types).</p>
-<p>The context can be used to store tensors that can be then retrieved
-during the backward pass.</p>
+<dt id="pytorch_tabnet.pretraining.TabNetPretrainer.fit">
+<code class="sig-name descname">fit</code><span class="sig-paren">(</span><em class="sig-param">X_train</em>, <em class="sig-param">eval_set=None</em>, <em class="sig-param">eval_name=None</em>, <em class="sig-param">loss_fn=None</em>, <em class="sig-param">pretraining_ratio=0.5</em>, <em class="sig-param">weights=0</em>, <em class="sig-param">max_epochs=100</em>, <em class="sig-param">patience=10</em>, <em class="sig-param">batch_size=1024</em>, <em class="sig-param">virtual_batch_size=128</em>, <em class="sig-param">num_workers=0</em>, <em class="sig-param">drop_last=True</em>, <em class="sig-param">callbacks=None</em>, <em class="sig-param">pin_memory=True</em>, <em class="sig-param">warm_start=False</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/pretraining.html#TabNetPretrainer.fit"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.pretraining.TabNetPretrainer.fit" title="Permalink to this definition">¶</a></dt>
+<dd><p>Train a neural network stored in self.network
+Using train_dataloader for training data and
+valid_dataloader for validation.</p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters</dt>
+<dd class="field-odd"><ul class="simple">
+<li><p><strong>X_train</strong> (<em>np.ndarray</em>) – Train set to reconstruct in self supervision</p></li>
+<li><p><strong>eval_set</strong> (<em>list of np.array</em>) – List of evaluation set
+The last one is used for early stopping</p></li>
+<li><p><strong>eval_name</strong> (<em>list of str</em>) – List of eval set names.</p></li>
+<li><p><strong>eval_metric</strong> (<em>list of str</em>) – List of evaluation metrics.
+The last metric is used for early stopping.</p></li>
+<li><p><strong>loss_fn</strong> (<em>callable</em><em> or </em><em>None</em>) – a PyTorch loss function
+should be left to None for self supervised and non experts</p></li>
+<li><p><strong>pretraining_ratio</strong> (<em>float</em>) – Between 0 and 1, percentage of feature to mask for reconstruction</p></li>
+<li><p><strong>weights</strong> (<em>np.array</em>) – Sampling weights for each example.</p></li>
+<li><p><strong>max_epochs</strong> (<em>int</em>) – Maximum number of epochs during training</p></li>
+<li><p><strong>patience</strong> (<em>int</em>) – Number of consecutive non improving epoch before early stopping</p></li>
+<li><p><strong>batch_size</strong> (<em>int</em>) – Training batch size</p></li>
+<li><p><strong>virtual_batch_size</strong> (<em>int</em>) – Batch size for Ghost Batch Normalization (virtual_batch_size &lt; batch_size)</p></li>
+<li><p><strong>num_workers</strong> (<em>int</em>) – Number of workers used in torch.utils.data.DataLoader</p></li>
+<li><p><strong>drop_last</strong> (<em>bool</em>) – Whether to drop last batch during training</p></li>
+<li><p><strong>callbacks</strong> (<em>list of callback function</em>) – List of custom callbacks</p></li>
+<li><p><strong>pin_memory</strong> (<em>bool</em>) – Whether to set pin_memory to True or False during training</p></li>
+</ul>
+</dd>
+</dl>
 </dd></dl>
 
+<dl class="attribute">
+<dt id="pytorch_tabnet.pretraining.TabNetPretrainer.optimizer_params">
+<code class="sig-name descname">optimizer_params</code><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.pretraining.TabNetPretrainer.optimizer_params" title="Permalink to this definition">¶</a></dt>
+<dd></dd></dl>
+
+<dl class="method">
+<dt id="pytorch_tabnet.pretraining.TabNetPretrainer.predict">
+<code class="sig-name descname">predict</code><span class="sig-paren">(</span><em class="sig-param">X</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/pretraining.html#TabNetPretrainer.predict"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.pretraining.TabNetPretrainer.predict" title="Permalink to this definition">¶</a></dt>
+<dd><p>Make predictions on a batch (valid)</p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters</dt>
+<dd class="field-odd"><p><strong>X</strong> (a :tensor: <cite>torch.Tensor</cite>) – Input data</p>
+</dd>
+<dt class="field-even">Returns</dt>
+<dd class="field-even"><p><strong>predictions</strong> – Predictions of the regression problem</p>
+</dd>
+<dt class="field-odd">Return type</dt>
+<dd class="field-odd"><p>np.array</p>
+</dd>
+</dl>
 </dd></dl>
 
-<dl class="class">
-<dt id="pytorch_tabnet.sparsemax.Entmoid15">
-<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.sparsemax.</code><code class="sig-name descname">Entmoid15</code><a class="reference internal" href="../_modules/pytorch_tabnet/sparsemax.html#Entmoid15"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.sparsemax.Entmoid15" title="Permalink to this definition">¶</a></dt>
-<dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">torch.autograd.function.Function</span></code></p>
-<p>A highly optimized equivalent of lambda x: Entmax15([x, 0])</p>
 <dl class="method">
-<dt id="pytorch_tabnet.sparsemax.Entmoid15.backward">
-<em class="property">static </em><code class="sig-name descname">backward</code><span class="sig-paren">(</span><em class="sig-param">ctx</em>, <em class="sig-param">grad_output</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/sparsemax.html#Entmoid15.backward"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.sparsemax.Entmoid15.backward" title="Permalink to this definition">¶</a></dt>
-<dd><p>Defines a formula for differentiating the operation.</p>
-<p>This function is to be overridden by all subclasses.</p>
-<p>It must accept a context <code class="xref py py-attr docutils literal notranslate"><span class="pre">ctx</span></code> as the first argument, followed by
-as many outputs did <a class="reference internal" href="#pytorch_tabnet.sparsemax.Entmoid15.forward" title="pytorch_tabnet.sparsemax.Entmoid15.forward"><code class="xref py py-func docutils literal notranslate"><span class="pre">forward()</span></code></a> return, and it should return as many
-tensors, as there were inputs to <a class="reference internal" href="#pytorch_tabnet.sparsemax.Entmoid15.forward" title="pytorch_tabnet.sparsemax.Entmoid15.forward"><code class="xref py py-func docutils literal notranslate"><span class="pre">forward()</span></code></a>. Each argument is the
-gradient w.r.t the given output, and each returned value should be the
-gradient w.r.t. the corresponding input.</p>
-<p>The context can be used to retrieve tensors saved during the forward
-pass. It also has an attribute <code class="xref py py-attr docutils literal notranslate"><span class="pre">ctx.needs_input_grad</span></code> as a tuple
-of booleans representing whether each input needs gradient. E.g.,
-<a class="reference internal" href="#pytorch_tabnet.sparsemax.Entmoid15.backward" title="pytorch_tabnet.sparsemax.Entmoid15.backward"><code class="xref py py-func docutils literal notranslate"><span class="pre">backward()</span></code></a> will have <code class="docutils literal notranslate"><span class="pre">ctx.needs_input_grad[0]</span> <span class="pre">=</span> <span class="pre">True</span></code> if the
-first input to <a class="reference internal" href="#pytorch_tabnet.sparsemax.Entmoid15.forward" title="pytorch_tabnet.sparsemax.Entmoid15.forward"><code class="xref py py-func docutils literal notranslate"><span class="pre">forward()</span></code></a> needs gradient computated w.r.t. the
-output.</p>
+<dt id="pytorch_tabnet.pretraining.TabNetPretrainer.prepare_target">
+<code class="sig-name descname">prepare_target</code><span class="sig-paren">(</span><em class="sig-param">y</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/pretraining.html#TabNetPretrainer.prepare_target"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.pretraining.TabNetPretrainer.prepare_target" title="Permalink to this definition">¶</a></dt>
+<dd><p>Prepare target before training.</p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters</dt>
+<dd class="field-odd"><p><strong>y</strong> (a :tensor: <cite>torch.Tensor</cite>) – Target matrix.</p>
+</dd>
+<dt class="field-even">Returns</dt>
+<dd class="field-even"><p>Converted target matrix.</p>
+</dd>
+<dt class="field-odd">Return type</dt>
+<dd class="field-odd"><p><cite>torch.Tensor</cite></p>
+</dd>
+</dl>
 </dd></dl>
 
+<dl class="attribute">
+<dt id="pytorch_tabnet.pretraining.TabNetPretrainer.scheduler_params">
+<code class="sig-name descname">scheduler_params</code><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.pretraining.TabNetPretrainer.scheduler_params" title="Permalink to this definition">¶</a></dt>
+<dd></dd></dl>
+
 <dl class="method">
-<dt id="pytorch_tabnet.sparsemax.Entmoid15.forward">
-<em class="property">static </em><code class="sig-name descname">forward</code><span class="sig-paren">(</span><em class="sig-param">ctx</em>, <em class="sig-param">input</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/sparsemax.html#Entmoid15.forward"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.sparsemax.Entmoid15.forward" title="Permalink to this definition">¶</a></dt>
-<dd><p>Performs the operation.</p>
-<p>This function is to be overridden by all subclasses.</p>
-<p>It must accept a context ctx as the first argument, followed by any
-number of arguments (tensors or other types).</p>
-<p>The context can be used to store tensors that can be then retrieved
-during the backward pass.</p>
+<dt id="pytorch_tabnet.pretraining.TabNetPretrainer.stack_batches">
+<code class="sig-name descname">stack_batches</code><span class="sig-paren">(</span><em class="sig-param">list_output</em>, <em class="sig-param">list_embedded_x</em>, <em class="sig-param">list_obfuscation</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/pretraining.html#TabNetPretrainer.stack_batches"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.pretraining.TabNetPretrainer.stack_batches" title="Permalink to this definition">¶</a></dt>
+<dd></dd></dl>
+
+<dl class="method">
+<dt id="pytorch_tabnet.pretraining.TabNetPretrainer.update_fit_params">
+<code class="sig-name descname">update_fit_params</code><span class="sig-paren">(</span><em class="sig-param">weights</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/pretraining.html#TabNetPretrainer.update_fit_params"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.pretraining.TabNetPretrainer.update_fit_params" title="Permalink to this definition">¶</a></dt>
+<dd><p>Set attributes relative to fit function.</p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters</dt>
+<dd class="field-odd"><ul class="simple">
+<li><p><strong>X_train</strong> (<em>np.ndarray</em>) – Train set</p></li>
+<li><p><strong>y_train</strong> (<em>np.array</em>) – Train targets</p></li>
+<li><p><strong>eval_set</strong> (<em>list of tuple</em>) – List of eval tuple set (X, y).</p></li>
+<li><p><strong>weights</strong> (<em>bool</em><em> or </em><em>dictionnary</em>) – 0 for no balancing
+1 for automated balancing</p></li>
+</ul>
+</dd>
+</dl>
+</dd></dl>
+
+</dd></dl>
+
+</section>
+<section id="module-pytorch_tabnet.utils">
+<span id="pytorch-tabnet-utils-module"></span><h2>pytorch_tabnet.utils module<a class="headerlink" href="#module-pytorch_tabnet.utils" title="Permalink to this headline">¶</a></h2>
+<dl class="class">
+<dt id="pytorch_tabnet.utils.ComplexEncoder">
+<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.utils.</code><code class="sig-name descname">ComplexEncoder</code><span class="sig-paren">(</span><em class="sig-param">*</em>, <em class="sig-param">skipkeys=False</em>, <em class="sig-param">ensure_ascii=True</em>, <em class="sig-param">check_circular=True</em>, <em class="sig-param">allow_nan=True</em>, <em class="sig-param">sort_keys=False</em>, <em class="sig-param">indent=None</em>, <em class="sig-param">separators=None</em>, <em class="sig-param">default=None</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/utils.html#ComplexEncoder"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.utils.ComplexEncoder" title="Permalink to this definition">¶</a></dt>
+<dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">json.encoder.JSONEncoder</span></code></p>
+<dl class="method">
+<dt id="pytorch_tabnet.utils.ComplexEncoder.default">
+<code class="sig-name descname">default</code><span class="sig-paren">(</span><em class="sig-param">obj</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/utils.html#ComplexEncoder.default"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.utils.ComplexEncoder.default" title="Permalink to this definition">¶</a></dt>
+<dd><p>Implement this method in a subclass such that it returns
+a serializable object for <code class="docutils literal notranslate"><span class="pre">o</span></code>, or calls the base implementation
+(to raise a <code class="docutils literal notranslate"><span class="pre">TypeError</span></code>).</p>
+<p>For example, to support arbitrary iterators, you could
+implement default like this:</p>
+<div class="highlight-default notranslate"><div class="highlight"><pre><span></span><span class="k">def</span> <span class="nf">default</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">o</span><span class="p">):</span>
+    <span class="k">try</span><span class="p">:</span>
+        <span class="n">iterable</span> <span class="o">=</span> <span class="nb">iter</span><span class="p">(</span><span class="n">o</span><span class="p">)</span>
+    <span class="k">except</span> <span class="ne">TypeError</span><span class="p">:</span>
+        <span class="k">pass</span>
+    <span class="k">else</span><span class="p">:</span>
+        <span class="k">return</span> <span class="nb">list</span><span class="p">(</span><span class="n">iterable</span><span class="p">)</span>
+    <span class="c1"># Let the base class default method raise the TypeError</span>
+    <span class="k">return</span> <span class="n">JSONEncoder</span><span class="o">.</span><span class="n">default</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">o</span><span class="p">)</span>
+</pre></div>
+</div>
+</dd></dl>
+
+</dd></dl>
+
+<dl class="class">
+<dt id="pytorch_tabnet.utils.PredictDataset">
+<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.utils.</code><code class="sig-name descname">PredictDataset</code><span class="sig-paren">(</span><em class="sig-param">x</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/utils.html#PredictDataset"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.utils.PredictDataset" title="Permalink to this definition">¶</a></dt>
+<dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">torch.utils.data.dataset.Dataset</span></code></p>
+<p>Format for numpy array</p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters</dt>
+<dd class="field-odd"><p><strong>X</strong> (<em>2D array</em>) – The input matrix</p>
+</dd>
+</dl>
+</dd></dl>
+
+<dl class="class">
+<dt id="pytorch_tabnet.utils.TorchDataset">
+<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.utils.</code><code class="sig-name descname">TorchDataset</code><span class="sig-paren">(</span><em class="sig-param">x</em>, <em class="sig-param">y</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/utils.html#TorchDataset"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.utils.TorchDataset" title="Permalink to this definition">¶</a></dt>
+<dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">torch.utils.data.dataset.Dataset</span></code></p>
+<p>Format for numpy array</p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters</dt>
+<dd class="field-odd"><ul class="simple">
+<li><p><strong>X</strong> (<em>2D array</em>) – The input matrix</p></li>
+<li><p><strong>y</strong> (<em>2D array</em>) – The one-hot encoded target</p></li>
+</ul>
+</dd>
+</dl>
+</dd></dl>
+
+<dl class="function">
+<dt id="pytorch_tabnet.utils.check_input">
+<code class="sig-prename descclassname">pytorch_tabnet.utils.</code><code class="sig-name descname">check_input</code><span class="sig-paren">(</span><em class="sig-param">X</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/utils.html#check_input"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.utils.check_input" title="Permalink to this definition">¶</a></dt>
+<dd><p>Raise a clear error if X is a pandas dataframe
+and check array according to scikit rules</p>
+</dd></dl>
+
+<dl class="function">
+<dt id="pytorch_tabnet.utils.check_warm_start">
+<code class="sig-prename descclassname">pytorch_tabnet.utils.</code><code class="sig-name descname">check_warm_start</code><span class="sig-paren">(</span><em class="sig-param">warm_start</em>, <em class="sig-param">from_unsupervised</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/utils.html#check_warm_start"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.utils.check_warm_start" title="Permalink to this definition">¶</a></dt>
+<dd><p>Gives a warning about ambiguous usage of the two parameters.</p>
+</dd></dl>
+
+<dl class="function">
+<dt id="pytorch_tabnet.utils.create_dataloaders">
+<code class="sig-prename descclassname">pytorch_tabnet.utils.</code><code class="sig-name descname">create_dataloaders</code><span class="sig-paren">(</span><em class="sig-param">X_train</em>, <em class="sig-param">y_train</em>, <em class="sig-param">eval_set</em>, <em class="sig-param">weights</em>, <em class="sig-param">batch_size</em>, <em class="sig-param">num_workers</em>, <em class="sig-param">drop_last</em>, <em class="sig-param">pin_memory</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/utils.html#create_dataloaders"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.utils.create_dataloaders" title="Permalink to this definition">¶</a></dt>
+<dd><p>Create dataloaders with or without subsampling depending on weights and balanced.</p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters</dt>
+<dd class="field-odd"><ul class="simple">
+<li><p><strong>X_train</strong> (<em>np.ndarray</em>) – Training data</p></li>
+<li><p><strong>y_train</strong> (<em>np.array</em>) – Mapped Training targets</p></li>
+<li><p><strong>eval_set</strong> (<em>list of tuple</em>) – List of eval tuple set (X, y)</p></li>
+<li><p><strong>weights</strong> (<em>either 0</em><em>, </em><em>1</em><em>, </em><em>dict</em><em> or </em><em>iterable</em>) – <p>if 0 (default) : no weights will be applied
+if 1 : classification only, will balanced class with inverse frequency
+if dict : keys are corresponding class values are sample weights
+if iterable : list or np array must be of length equal to nb elements</p>
+<blockquote>
+<div><p>in the training set</p>
+</div></blockquote>
+</p></li>
+<li><p><strong>batch_size</strong> (<em>int</em>) – how many samples per batch to load</p></li>
+<li><p><strong>num_workers</strong> (<em>int</em>) – how many subprocesses to use for data loading. 0 means that the data
+will be loaded in the main process</p></li>
+<li><p><strong>drop_last</strong> (<em>bool</em>) – set to True to drop the last incomplete batch, if the dataset size is not
+divisible by the batch size. If False and the size of dataset is not
+divisible by the batch size, then the last batch will be smaller</p></li>
+<li><p><strong>pin_memory</strong> (<em>bool</em>) – Whether to pin GPU memory during training</p></li>
+</ul>
+</dd>
+<dt class="field-even">Returns</dt>
+<dd class="field-even"><p><strong>train_dataloader, valid_dataloader</strong> – Training and validation dataloaders</p>
+</dd>
+<dt class="field-odd">Return type</dt>
+<dd class="field-odd"><p>torch.DataLoader, torch.DataLoader</p>
+</dd>
+</dl>
+</dd></dl>
+
+<dl class="function">
+<dt id="pytorch_tabnet.utils.create_explain_matrix">
+<code class="sig-prename descclassname">pytorch_tabnet.utils.</code><code class="sig-name descname">create_explain_matrix</code><span class="sig-paren">(</span><em class="sig-param">input_dim</em>, <em class="sig-param">cat_emb_dim</em>, <em class="sig-param">cat_idxs</em>, <em class="sig-param">post_embed_dim</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/utils.html#create_explain_matrix"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.utils.create_explain_matrix" title="Permalink to this definition">¶</a></dt>
+<dd><p>This is a computational trick.
+In order to rapidly sum importances from same embeddings
+to the initial index.</p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters</dt>
+<dd class="field-odd"><ul class="simple">
+<li><p><strong>input_dim</strong> (<em>int</em>) – Initial input dim</p></li>
+<li><p><strong>cat_emb_dim</strong> (<em>int</em><em> or </em><em>list of int</em>) – if int : size of embedding for all categorical feature
+if list of int : size of embedding for each categorical feature</p></li>
+<li><p><strong>cat_idxs</strong> (<em>list of int</em>) – Initial position of categorical features</p></li>
+<li><p><strong>post_embed_dim</strong> (<em>int</em>) – Post embedding inputs dimension</p></li>
+</ul>
+</dd>
+<dt class="field-even">Returns</dt>
+<dd class="field-even"><p><strong>reducing_matrix</strong> – Matrix of dim (post_embed_dim, input_dim)  to performe reduce</p>
+</dd>
+<dt class="field-odd">Return type</dt>
+<dd class="field-odd"><p>np.array</p>
+</dd>
+</dl>
+</dd></dl>
+
+<dl class="function">
+<dt id="pytorch_tabnet.utils.create_sampler">
+<code class="sig-prename descclassname">pytorch_tabnet.utils.</code><code class="sig-name descname">create_sampler</code><span class="sig-paren">(</span><em class="sig-param">weights</em>, <em class="sig-param">y_train</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/utils.html#create_sampler"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.utils.create_sampler" title="Permalink to this definition">¶</a></dt>
+<dd><p>This creates a sampler from the given weights</p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters</dt>
+<dd class="field-odd"><ul class="simple">
+<li><p><strong>weights</strong> (<em>either 0</em><em>, </em><em>1</em><em>, </em><em>dict</em><em> or </em><em>iterable</em>) – <p>if 0 (default) : no weights will be applied
+if 1 : classification only, will balanced class with inverse frequency
+if dict : keys are corresponding class values are sample weights
+if iterable : list or np array must be of length equal to nb elements</p>
+<blockquote>
+<div><p>in the training set</p>
+</div></blockquote>
+</p></li>
+<li><p><strong>y_train</strong> (<em>np.array</em>) – Training targets</p></li>
+</ul>
+</dd>
+</dl>
 </dd></dl>
 
+<dl class="function">
+<dt id="pytorch_tabnet.utils.define_device">
+<code class="sig-prename descclassname">pytorch_tabnet.utils.</code><code class="sig-name descname">define_device</code><span class="sig-paren">(</span><em class="sig-param">device_name</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/utils.html#define_device"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.utils.define_device" title="Permalink to this definition">¶</a></dt>
+<dd><p>Define the device to use during training and inference.
+If auto it will detect automatically whether to use cuda or cpu</p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters</dt>
+<dd class="field-odd"><p><strong>device_name</strong> (<em>str</em>) – Either “auto”, “cpu” or “cuda”</p>
+</dd>
+<dt class="field-even">Returns</dt>
+<dd class="field-even"><p>Either “cpu” or “cuda”</p>
+</dd>
+<dt class="field-odd">Return type</dt>
+<dd class="field-odd"><p>str</p>
+</dd>
+</dl>
 </dd></dl>
 
-<dl class="class">
-<dt id="pytorch_tabnet.sparsemax.Sparsemax">
-<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.sparsemax.</code><code class="sig-name descname">Sparsemax</code><span class="sig-paren">(</span><em class="sig-param">dim=-1</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/sparsemax.html#Sparsemax"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.sparsemax.Sparsemax" title="Permalink to this definition">¶</a></dt>
-<dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">torch.nn.modules.module.Module</span></code></p>
-<dl class="method">
-<dt id="pytorch_tabnet.sparsemax.Sparsemax.forward">
-<code class="sig-name descname">forward</code><span class="sig-paren">(</span><em class="sig-param">input</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/sparsemax.html#Sparsemax.forward"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.sparsemax.Sparsemax.forward" title="Permalink to this definition">¶</a></dt>
-<dd><p>Defines the computation performed at every call.</p>
-<p>Should be overridden by all subclasses.</p>
-<div class="admonition note">
-<p class="admonition-title">Note</p>
-<p>Although the recipe for forward pass needs to be defined within
-this function, one should call the <code class="xref py py-class docutils literal notranslate"><span class="pre">Module</span></code> instance afterwards
-instead of this since the former takes care of running the
-registered hooks while the latter silently ignores them.</p>
-</div>
+<dl class="function">
+<dt id="pytorch_tabnet.utils.filter_weights">
+<code class="sig-prename descclassname">pytorch_tabnet.utils.</code><code class="sig-name descname">filter_weights</code><span class="sig-paren">(</span><em class="sig-param">weights</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/utils.html#filter_weights"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.utils.filter_weights" title="Permalink to this definition">¶</a></dt>
+<dd><p>This function makes sure that weights are in correct format for
+regression and multitask TabNet</p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters</dt>
+<dd class="field-odd"><p><strong>weights</strong> (<em>int</em><em>, </em><em>dict</em><em> or </em><em>list</em>) – Initial weights parameters given by user</p>
+</dd>
+<dt class="field-even">Returns</dt>
+<dd class="field-even"><p><strong>None</strong></p>
+</dd>
+<dt class="field-odd">Return type</dt>
+<dd class="field-odd"><p>This function will only throw an error if format is wrong</p>
+</dd>
+</dl>
 </dd></dl>
 
+<dl class="function">
+<dt id="pytorch_tabnet.utils.validate_eval_set">
+<code class="sig-prename descclassname">pytorch_tabnet.utils.</code><code class="sig-name descname">validate_eval_set</code><span class="sig-paren">(</span><em class="sig-param">eval_set</em>, <em class="sig-param">eval_name</em>, <em class="sig-param">X_train</em>, <em class="sig-param">y_train</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/utils.html#validate_eval_set"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.utils.validate_eval_set" title="Permalink to this definition">¶</a></dt>
+<dd><p>Check if the shapes of eval_set are compatible with (X_train, y_train).</p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters</dt>
+<dd class="field-odd"><ul class="simple">
+<li><p><strong>eval_set</strong> (<em>list of tuple</em>) – List of eval tuple set (X, y).
+The last one is used for early stopping</p></li>
+<li><p><strong>eval_name</strong> (<em>list of str</em>) – List of eval set names.</p></li>
+<li><p><strong>X_train</strong> (<em>np.ndarray</em>) – Train owned products</p></li>
+<li><p><strong>y_train</strong> (<em>np.array</em>) – Train targeted products</p></li>
+</ul>
+</dd>
+<dt class="field-even">Returns</dt>
+<dd class="field-even"><p><ul class="simple">
+<li><p><strong>eval_names</strong> (<em>list of str</em>) – Validated list of eval_names.</p></li>
+<li><p><strong>eval_set</strong> (<em>list of tuple</em>) – Validated list of eval_set.</p></li>
+</ul>
+</p>
+</dd>
+</dl>
 </dd></dl>
 
+</section>
+<section id="module-pytorch_tabnet.multitask">
+<span id="pytorch-tabnet-multitask-module"></span><h2>pytorch_tabnet.multitask module<a class="headerlink" href="#module-pytorch_tabnet.multitask" title="Permalink to this headline">¶</a></h2>
 <dl class="class">
-<dt id="pytorch_tabnet.sparsemax.SparsemaxFunction">
-<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.sparsemax.</code><code class="sig-name descname">SparsemaxFunction</code><a class="reference internal" href="../_modules/pytorch_tabnet/sparsemax.html#SparsemaxFunction"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.sparsemax.SparsemaxFunction" title="Permalink to this definition">¶</a></dt>
-<dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">torch.autograd.function.Function</span></code></p>
-<p>An implementation of sparsemax (Martins &amp; Astudillo, 2016). See
-<a href="#id3"><span class="problematic" id="id4">:cite:`DBLP:journals/corr/MartinsA16`</span></a> for detailed description.
-By Ben Peters and Vlad Niculae</p>
-<dl class="method">
-<dt id="pytorch_tabnet.sparsemax.SparsemaxFunction.backward">
-<em class="property">static </em><code class="sig-name descname">backward</code><span class="sig-paren">(</span><em class="sig-param">ctx</em>, <em class="sig-param">grad_output</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/sparsemax.html#SparsemaxFunction.backward"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.sparsemax.SparsemaxFunction.backward" title="Permalink to this definition">¶</a></dt>
-<dd><p>Defines a formula for differentiating the operation.</p>
-<p>This function is to be overridden by all subclasses.</p>
-<p>It must accept a context <code class="xref py py-attr docutils literal notranslate"><span class="pre">ctx</span></code> as the first argument, followed by
-as many outputs did <a class="reference internal" href="#pytorch_tabnet.sparsemax.SparsemaxFunction.forward" title="pytorch_tabnet.sparsemax.SparsemaxFunction.forward"><code class="xref py py-func docutils literal notranslate"><span class="pre">forward()</span></code></a> return, and it should return as many
-tensors, as there were inputs to <a class="reference internal" href="#pytorch_tabnet.sparsemax.SparsemaxFunction.forward" title="pytorch_tabnet.sparsemax.SparsemaxFunction.forward"><code class="xref py py-func docutils literal notranslate"><span class="pre">forward()</span></code></a>. Each argument is the
-gradient w.r.t the given output, and each returned value should be the
-gradient w.r.t. the corresponding input.</p>
-<p>The context can be used to retrieve tensors saved during the forward
-pass. It also has an attribute <code class="xref py py-attr docutils literal notranslate"><span class="pre">ctx.needs_input_grad</span></code> as a tuple
-of booleans representing whether each input needs gradient. E.g.,
-<a class="reference internal" href="#pytorch_tabnet.sparsemax.SparsemaxFunction.backward" title="pytorch_tabnet.sparsemax.SparsemaxFunction.backward"><code class="xref py py-func docutils literal notranslate"><span class="pre">backward()</span></code></a> will have <code class="docutils literal notranslate"><span class="pre">ctx.needs_input_grad[0]</span> <span class="pre">=</span> <span class="pre">True</span></code> if the
-first input to <a class="reference internal" href="#pytorch_tabnet.sparsemax.SparsemaxFunction.forward" title="pytorch_tabnet.sparsemax.SparsemaxFunction.forward"><code class="xref py py-func docutils literal notranslate"><span class="pre">forward()</span></code></a> needs gradient computated w.r.t. the
-output.</p>
-</dd></dl>
+<dt id="pytorch_tabnet.multitask.TabNetMultiTaskClassifier">
+<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.multitask.</code><code class="sig-name descname">TabNetMultiTaskClassifier</code><span class="sig-paren">(</span><em class="sig-param">n_d: int = 8</em>, <em class="sig-param">n_a: int = 8</em>, <em class="sig-param">n_steps: int = 3</em>, <em class="sig-param">gamma: float = 1.3</em>, <em class="sig-param">cat_idxs: List[int] = &lt;factory&gt;</em>, <em class="sig-param">cat_dims: List[int] = &lt;factory&gt;</em>, <em class="sig-param">cat_emb_dim: int = 1</em>, <em class="sig-param">n_independent: int = 2</em>, <em class="sig-param">n_shared: int = 2</em>, <em class="sig-param">epsilon: float = 1e-15</em>, <em class="sig-param">momentum: float = 0.02</em>, <em class="sig-param">lambda_sparse: float = 0.001</em>, <em class="sig-param">seed: int = 0</em>, <em class="sig-param">clip_value: int = 1</em>, <em class="sig-param">verbose: int = 1</em>, <em class="sig-param">optimizer_fn: Any = &lt;class 'torch.optim.adam.Adam'&gt;</em>, <em class="sig-param">optimizer_params: Dict = &lt;factory&gt;</em>, <em class="sig-param">scheduler_fn: Any = None</em>, <em class="sig-param">scheduler_params: Dict = &lt;factory&gt;</em>, <em class="sig-param">mask_type: str = 'sparsemax'</em>, <em class="sig-param">input_dim: int = None</em>, <em class="sig-param">output_dim: int = None</em>, <em class="sig-param">device_name: str = 'auto'</em>, <em class="sig-param">n_shared_decoder: int = 1</em>, <em class="sig-param">n_indep_decoder: int = 1</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/multitask.html#TabNetMultiTaskClassifier"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.multitask.TabNetMultiTaskClassifier" title="Permalink to this definition">¶</a></dt>
+<dd><p>Bases: <a class="reference internal" href="#pytorch_tabnet.abstract_model.TabModel" title="pytorch_tabnet.abstract_model.TabModel"><code class="xref py py-class docutils literal notranslate"><span class="pre">pytorch_tabnet.abstract_model.TabModel</span></code></a></p>
+<dl class="attribute">
+<dt id="pytorch_tabnet.multitask.TabNetMultiTaskClassifier.cat_dims">
+<code class="sig-name descname">cat_dims</code><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.multitask.TabNetMultiTaskClassifier.cat_dims" title="Permalink to this definition">¶</a></dt>
+<dd></dd></dl>
+
+<dl class="attribute">
+<dt id="pytorch_tabnet.multitask.TabNetMultiTaskClassifier.cat_idxs">
+<code class="sig-name descname">cat_idxs</code><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.multitask.TabNetMultiTaskClassifier.cat_idxs" title="Permalink to this definition">¶</a></dt>
+<dd></dd></dl>
 
 <dl class="method">
-<dt id="pytorch_tabnet.sparsemax.SparsemaxFunction.forward">
-<em class="property">static </em><code class="sig-name descname">forward</code><span class="sig-paren">(</span><em class="sig-param">ctx</em>, <em class="sig-param">input</em>, <em class="sig-param">dim=-1</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/sparsemax.html#SparsemaxFunction.forward"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.sparsemax.SparsemaxFunction.forward" title="Permalink to this definition">¶</a></dt>
-<dd><p>sparsemax: normalizing sparse transform (a la softmax)</p>
+<dt id="pytorch_tabnet.multitask.TabNetMultiTaskClassifier.compute_loss">
+<code class="sig-name descname">compute_loss</code><span class="sig-paren">(</span><em class="sig-param">y_pred</em>, <em class="sig-param">y_true</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/multitask.html#TabNetMultiTaskClassifier.compute_loss"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.multitask.TabNetMultiTaskClassifier.compute_loss" title="Permalink to this definition">¶</a></dt>
+<dd><p>Computes the loss according to network output and targets</p>
 <dl class="field-list simple">
 <dt class="field-odd">Parameters</dt>
 <dd class="field-odd"><ul class="simple">
-<li><p><strong>ctx</strong> (<em>torch.autograd.function._ContextMethodMixin</em>) – </p></li>
-<li><p><strong>input</strong> (<em>torch.Tensor</em>) – any shape</p></li>
-<li><p><strong>dim</strong> (<em>int</em>) – dimension along which to apply sparsemax</p></li>
+<li><p><strong>y_pred</strong> (<em>list of tensors</em>) – Output of network</p></li>
+<li><p><strong>y_true</strong> (<em>LongTensor</em>) – Targets label encoded</p></li>
 </ul>
 </dd>
 <dt class="field-even">Returns</dt>
-<dd class="field-even"><p><strong>output</strong> – same shape as input</p>
+<dd class="field-even"><p><strong>loss</strong> – output of loss function(s)</p>
 </dd>
 <dt class="field-odd">Return type</dt>
 <dd class="field-odd"><p>torch.Tensor</p>
@@ -2112,216 +2075,319 @@ <h3>Multi-class / multi-label utility function<a class="headerlink" href="#multi
 </dl>
 </dd></dl>
 
-</dd></dl>
-
-<dl class="function">
-<dt id="pytorch_tabnet.sparsemax.entmax15">
-<code class="sig-prename descclassname">pytorch_tabnet.sparsemax.</code><code class="sig-name descname">entmax15</code><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="headerlink" href="#pytorch_tabnet.sparsemax.entmax15" title="Permalink to this definition">¶</a></dt>
-<dd></dd></dl>
-
-<dl class="function">
-<dt id="pytorch_tabnet.sparsemax.entmoid15">
-<code class="sig-prename descclassname">pytorch_tabnet.sparsemax.</code><code class="sig-name descname">entmoid15</code><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="headerlink" href="#pytorch_tabnet.sparsemax.entmoid15" title="Permalink to this definition">¶</a></dt>
-<dd></dd></dl>
-
-<dl class="function">
-<dt id="pytorch_tabnet.sparsemax.sparsemax">
-<code class="sig-prename descclassname">pytorch_tabnet.sparsemax.</code><code class="sig-name descname">sparsemax</code><span class="sig-paren">(</span><span class="sig-paren">)</span><a class="headerlink" href="#pytorch_tabnet.sparsemax.sparsemax" title="Permalink to this definition">¶</a></dt>
+<dl class="attribute">
+<dt id="pytorch_tabnet.multitask.TabNetMultiTaskClassifier.optimizer_params">
+<code class="sig-name descname">optimizer_params</code><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.multitask.TabNetMultiTaskClassifier.optimizer_params" title="Permalink to this definition">¶</a></dt>
 <dd></dd></dl>
 
-</div>
-<div class="section" id="module-pytorch_tabnet.metrics">
-<span id="pytorch-tabnet-metrics-module"></span><h2>pytorch_tabnet.metrics module<a class="headerlink" href="#module-pytorch_tabnet.metrics" title="Permalink to this headline">¶</a></h2>
-<dl class="class">
-<dt id="pytorch_tabnet.metrics.AUC">
-<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.metrics.</code><code class="sig-name descname">AUC</code><a class="reference internal" href="../_modules/pytorch_tabnet/metrics.html#AUC"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.metrics.AUC" title="Permalink to this definition">¶</a></dt>
-<dd><p>Bases: <a class="reference internal" href="#pytorch_tabnet.metrics.Metric" title="pytorch_tabnet.metrics.Metric"><code class="xref py py-class docutils literal notranslate"><span class="pre">pytorch_tabnet.metrics.Metric</span></code></a></p>
-<p>AUC.</p>
+<dl class="method">
+<dt id="pytorch_tabnet.multitask.TabNetMultiTaskClassifier.predict">
+<code class="sig-name descname">predict</code><span class="sig-paren">(</span><em class="sig-param">X</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/multitask.html#TabNetMultiTaskClassifier.predict"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.multitask.TabNetMultiTaskClassifier.predict" title="Permalink to this definition">¶</a></dt>
+<dd><p>Make predictions on a batch (valid)</p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters</dt>
+<dd class="field-odd"><p><strong>X</strong> (a :tensor: <cite>torch.Tensor</cite>) – Input data</p>
+</dd>
+<dt class="field-even">Returns</dt>
+<dd class="field-even"><p><strong>results</strong> – Predictions of the most probable class</p>
+</dd>
+<dt class="field-odd">Return type</dt>
+<dd class="field-odd"><p>np.array</p>
+</dd>
+</dl>
 </dd></dl>
 
-<dl class="class">
-<dt id="pytorch_tabnet.metrics.Accuracy">
-<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.metrics.</code><code class="sig-name descname">Accuracy</code><a class="reference internal" href="../_modules/pytorch_tabnet/metrics.html#Accuracy"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.metrics.Accuracy" title="Permalink to this definition">¶</a></dt>
-<dd><p>Bases: <a class="reference internal" href="#pytorch_tabnet.metrics.Metric" title="pytorch_tabnet.metrics.Metric"><code class="xref py py-class docutils literal notranslate"><span class="pre">pytorch_tabnet.metrics.Metric</span></code></a></p>
-<p>Accuracy.</p>
+<dl class="method">
+<dt id="pytorch_tabnet.multitask.TabNetMultiTaskClassifier.predict_proba">
+<code class="sig-name descname">predict_proba</code><span class="sig-paren">(</span><em class="sig-param">X</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/multitask.html#TabNetMultiTaskClassifier.predict_proba"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.multitask.TabNetMultiTaskClassifier.predict_proba" title="Permalink to this definition">¶</a></dt>
+<dd><p>Make predictions for classification on a batch (valid)</p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters</dt>
+<dd class="field-odd"><p><strong>X</strong> (a :tensor: <cite>torch.Tensor</cite>) – Input data</p>
+</dd>
+<dt class="field-even">Returns</dt>
+<dd class="field-even"><p><strong>res</strong></p>
+</dd>
+<dt class="field-odd">Return type</dt>
+<dd class="field-odd"><p>list of np.ndarray</p>
+</dd>
+</dl>
 </dd></dl>
 
-<dl class="class">
-<dt id="pytorch_tabnet.metrics.BalancedAccuracy">
-<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.metrics.</code><code class="sig-name descname">BalancedAccuracy</code><a class="reference internal" href="../_modules/pytorch_tabnet/metrics.html#BalancedAccuracy"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.metrics.BalancedAccuracy" title="Permalink to this definition">¶</a></dt>
-<dd><p>Bases: <a class="reference internal" href="#pytorch_tabnet.metrics.Metric" title="pytorch_tabnet.metrics.Metric"><code class="xref py py-class docutils literal notranslate"><span class="pre">pytorch_tabnet.metrics.Metric</span></code></a></p>
-<p>Balanced Accuracy.</p>
+<dl class="method">
+<dt id="pytorch_tabnet.multitask.TabNetMultiTaskClassifier.prepare_target">
+<code class="sig-name descname">prepare_target</code><span class="sig-paren">(</span><em class="sig-param">y</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/multitask.html#TabNetMultiTaskClassifier.prepare_target"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.multitask.TabNetMultiTaskClassifier.prepare_target" title="Permalink to this definition">¶</a></dt>
+<dd><p>Prepare target before training.</p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters</dt>
+<dd class="field-odd"><p><strong>y</strong> (a :tensor: <cite>torch.Tensor</cite>) – Target matrix.</p>
+</dd>
+<dt class="field-even">Returns</dt>
+<dd class="field-even"><p>Converted target matrix.</p>
+</dd>
+<dt class="field-odd">Return type</dt>
+<dd class="field-odd"><p><cite>torch.Tensor</cite></p>
+</dd>
+</dl>
 </dd></dl>
 
-<dl class="class">
-<dt id="pytorch_tabnet.metrics.LogLoss">
-<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.metrics.</code><code class="sig-name descname">LogLoss</code><a class="reference internal" href="../_modules/pytorch_tabnet/metrics.html#LogLoss"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.metrics.LogLoss" title="Permalink to this definition">¶</a></dt>
-<dd><p>Bases: <a class="reference internal" href="#pytorch_tabnet.metrics.Metric" title="pytorch_tabnet.metrics.Metric"><code class="xref py py-class docutils literal notranslate"><span class="pre">pytorch_tabnet.metrics.Metric</span></code></a></p>
-<p>LogLoss.</p>
-</dd></dl>
+<dl class="attribute">
+<dt id="pytorch_tabnet.multitask.TabNetMultiTaskClassifier.scheduler_params">
+<code class="sig-name descname">scheduler_params</code><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.multitask.TabNetMultiTaskClassifier.scheduler_params" title="Permalink to this definition">¶</a></dt>
+<dd></dd></dl>
 
-<dl class="class">
-<dt id="pytorch_tabnet.metrics.MAE">
-<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.metrics.</code><code class="sig-name descname">MAE</code><a class="reference internal" href="../_modules/pytorch_tabnet/metrics.html#MAE"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.metrics.MAE" title="Permalink to this definition">¶</a></dt>
-<dd><p>Bases: <a class="reference internal" href="#pytorch_tabnet.metrics.Metric" title="pytorch_tabnet.metrics.Metric"><code class="xref py py-class docutils literal notranslate"><span class="pre">pytorch_tabnet.metrics.Metric</span></code></a></p>
-<p>Mean Absolute Error.</p>
+<dl class="method">
+<dt id="pytorch_tabnet.multitask.TabNetMultiTaskClassifier.stack_batches">
+<code class="sig-name descname">stack_batches</code><span class="sig-paren">(</span><em class="sig-param">list_y_true</em>, <em class="sig-param">list_y_score</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/multitask.html#TabNetMultiTaskClassifier.stack_batches"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.multitask.TabNetMultiTaskClassifier.stack_batches" title="Permalink to this definition">¶</a></dt>
+<dd></dd></dl>
+
+<dl class="method">
+<dt id="pytorch_tabnet.multitask.TabNetMultiTaskClassifier.update_fit_params">
+<code class="sig-name descname">update_fit_params</code><span class="sig-paren">(</span><em class="sig-param">X_train</em>, <em class="sig-param">y_train</em>, <em class="sig-param">eval_set</em>, <em class="sig-param">weights</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/multitask.html#TabNetMultiTaskClassifier.update_fit_params"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.multitask.TabNetMultiTaskClassifier.update_fit_params" title="Permalink to this definition">¶</a></dt>
+<dd><p>Set attributes relative to fit function.</p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters</dt>
+<dd class="field-odd"><ul class="simple">
+<li><p><strong>X_train</strong> (<em>np.ndarray</em>) – Train set</p></li>
+<li><p><strong>y_train</strong> (<em>np.array</em>) – Train targets</p></li>
+<li><p><strong>eval_set</strong> (<em>list of tuple</em>) – List of eval tuple set (X, y).</p></li>
+<li><p><strong>weights</strong> (<em>bool</em><em> or </em><em>dictionnary</em>) – 0 for no balancing
+1 for automated balancing</p></li>
+</ul>
+</dd>
+</dl>
 </dd></dl>
 
-<dl class="class">
-<dt id="pytorch_tabnet.metrics.MSE">
-<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.metrics.</code><code class="sig-name descname">MSE</code><a class="reference internal" href="../_modules/pytorch_tabnet/metrics.html#MSE"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.metrics.MSE" title="Permalink to this definition">¶</a></dt>
-<dd><p>Bases: <a class="reference internal" href="#pytorch_tabnet.metrics.Metric" title="pytorch_tabnet.metrics.Metric"><code class="xref py py-class docutils literal notranslate"><span class="pre">pytorch_tabnet.metrics.Metric</span></code></a></p>
-<p>Mean Squared Error.</p>
 </dd></dl>
 
-<dl class="class">
-<dt id="pytorch_tabnet.metrics.Metric">
-<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.metrics.</code><code class="sig-name descname">Metric</code><a class="reference internal" href="../_modules/pytorch_tabnet/metrics.html#Metric"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.metrics.Metric" title="Permalink to this definition">¶</a></dt>
-<dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">object</span></code></p>
-<dl class="method">
-<dt id="pytorch_tabnet.metrics.Metric.get_metrics_by_names">
-<em class="property">classmethod </em><code class="sig-name descname">get_metrics_by_names</code><span class="sig-paren">(</span><em class="sig-param">names</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/metrics.html#Metric.get_metrics_by_names"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.metrics.Metric.get_metrics_by_names" title="Permalink to this definition">¶</a></dt>
-<dd><p>Get list of metric classes.</p>
+</section>
+<section id="module-pytorch_tabnet.multiclass_utils">
+<span id="pytorch-tabnet-multiclass-utils-module"></span><h2>pytorch_tabnet.multiclass_utils module<a class="headerlink" href="#module-pytorch_tabnet.multiclass_utils" title="Permalink to this headline">¶</a></h2>
+<section id="multi-class-multi-label-utility-function">
+<h3>Multi-class / multi-label utility function<a class="headerlink" href="#multi-class-multi-label-utility-function" title="Permalink to this headline">¶</a></h3>
+<dl class="function">
+<dt id="pytorch_tabnet.multiclass_utils.assert_all_finite">
+<code class="sig-prename descclassname">pytorch_tabnet.multiclass_utils.</code><code class="sig-name descname">assert_all_finite</code><span class="sig-paren">(</span><em class="sig-param">X</em>, <em class="sig-param">allow_nan=False</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/multiclass_utils.html#assert_all_finite"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.multiclass_utils.assert_all_finite" title="Permalink to this definition">¶</a></dt>
+<dd><p>Throw a ValueError if X contains NaN or infinity.</p>
 <dl class="field-list simple">
 <dt class="field-odd">Parameters</dt>
 <dd class="field-odd"><ul class="simple">
-<li><p><strong>cls</strong> (<a class="reference internal" href="#pytorch_tabnet.metrics.Metric" title="pytorch_tabnet.metrics.Metric"><em>Metric</em></a>) – Metric class.</p></li>
-<li><p><strong>names</strong> (<em>list</em>) – List of metric names.</p></li>
+<li><p><strong>X</strong> (<em>array</em><em> or </em><em>sparse matrix</em>) – </p></li>
+<li><p><strong>allow_nan</strong> (<em>bool</em>) – </p></li>
 </ul>
 </dd>
-<dt class="field-even">Returns</dt>
-<dd class="field-even"><p><strong>metrics</strong> – List of metric classes.</p>
-</dd>
-<dt class="field-odd">Return type</dt>
-<dd class="field-odd"><p>list</p>
-</dd>
 </dl>
 </dd></dl>
 
-</dd></dl>
-
-<dl class="class">
-<dt id="pytorch_tabnet.metrics.MetricContainer">
-<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.metrics.</code><code class="sig-name descname">MetricContainer</code><span class="sig-paren">(</span><em class="sig-param">metric_names: List[str], prefix: str = ''</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/metrics.html#MetricContainer"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.metrics.MetricContainer" title="Permalink to this definition">¶</a></dt>
-<dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">object</span></code></p>
-<p>Container holding a list of metrics.</p>
+<dl class="function">
+<dt id="pytorch_tabnet.multiclass_utils.check_classification_targets">
+<code class="sig-prename descclassname">pytorch_tabnet.multiclass_utils.</code><code class="sig-name descname">check_classification_targets</code><span class="sig-paren">(</span><em class="sig-param">y</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/multiclass_utils.html#check_classification_targets"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.multiclass_utils.check_classification_targets" title="Permalink to this definition">¶</a></dt>
+<dd><p>Ensure that target y is of a non-regression type.</p>
+<dl class="simple">
+<dt>Only the following target types (as defined in type_of_target) are allowed:</dt><dd><p>‘binary’, ‘multiclass’, ‘multiclass-multioutput’,
+‘multilabel-indicator’, ‘multilabel-sequences’</p>
+</dd>
+</dl>
 <dl class="field-list simple">
 <dt class="field-odd">Parameters</dt>
-<dd class="field-odd"><ul class="simple">
-<li><p><strong>metric_names</strong> (<em>list of str</em>) – List of metric names.</p></li>
-<li><p><strong>prefix</strong> (<em>str</em>) – Prefix of metric names.</p></li>
-</ul>
+<dd class="field-odd"><p><strong>y</strong> (<em>array-like</em>) – </p>
 </dd>
 </dl>
-<dl class="attribute">
-<dt id="pytorch_tabnet.metrics.MetricContainer.metric_names">
-<code class="sig-name descname">metric_names</code><em class="property">: List[str]</em><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.metrics.MetricContainer.metric_names" title="Permalink to this definition">¶</a></dt>
-<dd></dd></dl>
+</dd></dl>
 
-<dl class="attribute">
-<dt id="pytorch_tabnet.metrics.MetricContainer.prefix">
-<code class="sig-name descname">prefix</code><em class="property">: str</em><em class="property"> = ''</em><a class="headerlink" href="#pytorch_tabnet.metrics.MetricContainer.prefix" title="Permalink to this definition">¶</a></dt>
+<dl class="function">
+<dt id="pytorch_tabnet.multiclass_utils.check_output_dim">
+<code class="sig-prename descclassname">pytorch_tabnet.multiclass_utils.</code><code class="sig-name descname">check_output_dim</code><span class="sig-paren">(</span><em class="sig-param">labels</em>, <em class="sig-param">y</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/multiclass_utils.html#check_output_dim"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.multiclass_utils.check_output_dim" title="Permalink to this definition">¶</a></dt>
 <dd></dd></dl>
 
-</dd></dl>
-
-<dl class="class">
-<dt id="pytorch_tabnet.metrics.RMSE">
-<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.metrics.</code><code class="sig-name descname">RMSE</code><a class="reference internal" href="../_modules/pytorch_tabnet/metrics.html#RMSE"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.metrics.RMSE" title="Permalink to this definition">¶</a></dt>
-<dd><p>Bases: <a class="reference internal" href="#pytorch_tabnet.metrics.Metric" title="pytorch_tabnet.metrics.Metric"><code class="xref py py-class docutils literal notranslate"><span class="pre">pytorch_tabnet.metrics.Metric</span></code></a></p>
-<p>Root Mean Squared Error.</p>
-</dd></dl>
+<dl class="function">
+<dt id="pytorch_tabnet.multiclass_utils.check_unique_type">
+<code class="sig-prename descclassname">pytorch_tabnet.multiclass_utils.</code><code class="sig-name descname">check_unique_type</code><span class="sig-paren">(</span><em class="sig-param">y</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/multiclass_utils.html#check_unique_type"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.multiclass_utils.check_unique_type" title="Permalink to this definition">¶</a></dt>
+<dd></dd></dl>
 
-<dl class="class">
-<dt id="pytorch_tabnet.metrics.RMSLE">
-<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.metrics.</code><code class="sig-name descname">RMSLE</code><a class="reference internal" href="../_modules/pytorch_tabnet/metrics.html#RMSLE"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.metrics.RMSLE" title="Permalink to this definition">¶</a></dt>
-<dd><p>Bases: <a class="reference internal" href="#pytorch_tabnet.metrics.Metric" title="pytorch_tabnet.metrics.Metric"><code class="xref py py-class docutils literal notranslate"><span class="pre">pytorch_tabnet.metrics.Metric</span></code></a></p>
-<p>Mean squared logarithmic error regression loss.
-Scikit-implementation:
-<a class="reference external" href="https://scikit-learn.org/stable/modules/generated/sklearn.metrics.mean_squared_log_error.html">https://scikit-learn.org/stable/modules/generated/sklearn.metrics.mean_squared_log_error.html</a>
-Note: In order to avoid error, negative predictions are clipped to 0.
-This means that you should clip negative predictions manually after calling predict.</p>
+<dl class="function">
+<dt id="pytorch_tabnet.multiclass_utils.infer_multitask_output">
+<code class="sig-prename descclassname">pytorch_tabnet.multiclass_utils.</code><code class="sig-name descname">infer_multitask_output</code><span class="sig-paren">(</span><em class="sig-param">y_train</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/multiclass_utils.html#infer_multitask_output"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.multiclass_utils.infer_multitask_output" title="Permalink to this definition">¶</a></dt>
+<dd><p>Infer output_dim from targets
+This is for multiple tasks.</p>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters</dt>
+<dd class="field-odd"><p><strong>y_train</strong> (<em>np.ndarray</em>) – Training targets</p>
+</dd>
+<dt class="field-even">Returns</dt>
+<dd class="field-even"><p><ul class="simple">
+<li><p><strong>tasks_dims</strong> (<em>list</em>) – Number of classes for output</p></li>
+<li><p><strong>tasks_labels</strong> (<em>list</em>) – List of sorted list of initial classes</p></li>
+</ul>
+</p>
+</dd>
+</dl>
 </dd></dl>
 
-<dl class="class">
-<dt id="pytorch_tabnet.metrics.UnsupMetricContainer">
-<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.metrics.</code><code class="sig-name descname">UnsupMetricContainer</code><span class="sig-paren">(</span><em class="sig-param">metric_names: List[str], prefix: str = ''</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/metrics.html#UnsupMetricContainer"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.metrics.UnsupMetricContainer" title="Permalink to this definition">¶</a></dt>
-<dd><p>Bases: <code class="xref py py-class docutils literal notranslate"><span class="pre">object</span></code></p>
-<p>Container holding a list of metrics.</p>
+<dl class="function">
+<dt id="pytorch_tabnet.multiclass_utils.infer_output_dim">
+<code class="sig-prename descclassname">pytorch_tabnet.multiclass_utils.</code><code class="sig-name descname">infer_output_dim</code><span class="sig-paren">(</span><em class="sig-param">y_train</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/multiclass_utils.html#infer_output_dim"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.multiclass_utils.infer_output_dim" title="Permalink to this definition">¶</a></dt>
+<dd><p>Infer output_dim from targets</p>
 <dl class="field-list simple">
 <dt class="field-odd">Parameters</dt>
-<dd class="field-odd"><ul class="simple">
-<li><p><strong>y_pred</strong> (<em>torch.Tensor</em><em> or </em><em>np.array</em>) – Reconstructed prediction (with embeddings)</p></li>
-<li><p><strong>embedded_x</strong> (<em>torch.Tensor</em>) – Original input embedded by network</p></li>
-<li><p><strong>obf_vars</strong> (<em>torch.Tensor</em>) – Binary mask for obfuscated variables.
-1 means the variables was obfuscated so reconstruction is based on this.</p></li>
+<dd class="field-odd"><p><strong>y_train</strong> (<em>np.array</em>) – Training targets</p>
+</dd>
+<dt class="field-even">Returns</dt>
+<dd class="field-even"><p><ul class="simple">
+<li><p><strong>output_dim</strong> (<em>int</em>) – Number of classes for output</p></li>
+<li><p><strong>train_labels</strong> (<em>list</em>) – Sorted list of initial classes</p></li>
 </ul>
+</p>
 </dd>
 </dl>
-<dl class="attribute">
-<dt id="pytorch_tabnet.metrics.UnsupMetricContainer.metric_names">
-<code class="sig-name descname">metric_names</code><em class="property">: List[str]</em><em class="property"> = None</em><a class="headerlink" href="#pytorch_tabnet.metrics.UnsupMetricContainer.metric_names" title="Permalink to this definition">¶</a></dt>
-<dd></dd></dl>
-
-<dl class="attribute">
-<dt id="pytorch_tabnet.metrics.UnsupMetricContainer.prefix">
-<code class="sig-name descname">prefix</code><em class="property">: str</em><em class="property"> = ''</em><a class="headerlink" href="#pytorch_tabnet.metrics.UnsupMetricContainer.prefix" title="Permalink to this definition">¶</a></dt>
-<dd></dd></dl>
-
 </dd></dl>
 
 <dl class="function">
-<dt id="pytorch_tabnet.metrics.UnsupervisedLoss">
-<code class="sig-prename descclassname">pytorch_tabnet.metrics.</code><code class="sig-name descname">UnsupervisedLoss</code><span class="sig-paren">(</span><em class="sig-param">y_pred</em>, <em class="sig-param">embedded_x</em>, <em class="sig-param">obf_vars</em>, <em class="sig-param">eps=1e-09</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/metrics.html#UnsupervisedLoss"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.metrics.UnsupervisedLoss" title="Permalink to this definition">¶</a></dt>
-<dd><p>Implements unsupervised loss function.
-This differs from orginal paper as it’s scaled to be batch size independent
-and number of features reconstructed independent (by taking the mean)</p>
+<dt id="pytorch_tabnet.multiclass_utils.is_multilabel">
+<code class="sig-prename descclassname">pytorch_tabnet.multiclass_utils.</code><code class="sig-name descname">is_multilabel</code><span class="sig-paren">(</span><em class="sig-param">y</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/multiclass_utils.html#is_multilabel"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.multiclass_utils.is_multilabel" title="Permalink to this definition">¶</a></dt>
+<dd><p>Check if <code class="docutils literal notranslate"><span class="pre">y</span></code> is in a multilabel format.</p>
 <dl class="field-list simple">
 <dt class="field-odd">Parameters</dt>
-<dd class="field-odd"><ul class="simple">
-<li><p><strong>y_pred</strong> (<em>torch.Tensor</em><em> or </em><em>np.array</em>) – Reconstructed prediction (with embeddings)</p></li>
-<li><p><strong>embedded_x</strong> (<em>torch.Tensor</em>) – Original input embedded by network</p></li>
-<li><p><strong>obf_vars</strong> (<em>torch.Tensor</em>) – Binary mask for obfuscated variables.
-1 means the variable was obfuscated so reconstruction is based on this.</p></li>
-<li><p><strong>eps</strong> (<em>float</em>) – A small floating point to avoid ZeroDivisionError
-This can happen in degenerated case when a feature has only one value</p></li>
-</ul>
+<dd class="field-odd"><p><strong>y</strong> (<em>numpy array of shape</em><em> [</em><em>n_samples</em><em>]</em>) – Target values.</p>
 </dd>
 <dt class="field-even">Returns</dt>
-<dd class="field-even"><p><strong>loss</strong> – Unsupervised loss, average value over batch samples.</p>
+<dd class="field-even"><p><strong>out</strong> – Return <code class="docutils literal notranslate"><span class="pre">True</span></code>, if <code class="docutils literal notranslate"><span class="pre">y</span></code> is in a multilabel format, else <code class="docutils literal notranslate"><span class="pre">`False</span></code>.</p>
 </dd>
 <dt class="field-odd">Return type</dt>
-<dd class="field-odd"><p>torch float</p>
+<dd class="field-odd"><p>bool</p>
 </dd>
 </dl>
+<p class="rubric">Examples</p>
+<div class="doctest highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
+<span class="gp">&gt;&gt;&gt; </span><span class="kn">from</span> <span class="nn">sklearn.utils.multiclass</span> <span class="kn">import</span> <span class="n">is_multilabel</span>
+<span class="gp">&gt;&gt;&gt; </span><span class="n">is_multilabel</span><span class="p">([</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">])</span>
+<span class="go">False</span>
+<span class="gp">&gt;&gt;&gt; </span><span class="n">is_multilabel</span><span class="p">([[</span><span class="mi">1</span><span class="p">],</span> <span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">2</span><span class="p">],</span> <span class="p">[]])</span>
+<span class="go">False</span>
+<span class="gp">&gt;&gt;&gt; </span><span class="n">is_multilabel</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">([[</span><span class="mi">1</span><span class="p">,</span> <span class="mi">0</span><span class="p">],</span> <span class="p">[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">]]))</span>
+<span class="go">True</span>
+<span class="gp">&gt;&gt;&gt; </span><span class="n">is_multilabel</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">([[</span><span class="mi">1</span><span class="p">],</span> <span class="p">[</span><span class="mi">0</span><span class="p">],</span> <span class="p">[</span><span class="mi">0</span><span class="p">]]))</span>
+<span class="go">False</span>
+<span class="gp">&gt;&gt;&gt; </span><span class="n">is_multilabel</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">([[</span><span class="mi">1</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">0</span><span class="p">]]))</span>
+<span class="go">True</span>
+</pre></div>
+</div>
 </dd></dl>
 
-<dl class="class">
-<dt id="pytorch_tabnet.metrics.UnsupervisedMetric">
-<em class="property">class </em><code class="sig-prename descclassname">pytorch_tabnet.metrics.</code><code class="sig-name descname">UnsupervisedMetric</code><a class="reference internal" href="../_modules/pytorch_tabnet/metrics.html#UnsupervisedMetric"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.metrics.UnsupervisedMetric" title="Permalink to this definition">¶</a></dt>
-<dd><p>Bases: <a class="reference internal" href="#pytorch_tabnet.metrics.Metric" title="pytorch_tabnet.metrics.Metric"><code class="xref py py-class docutils literal notranslate"><span class="pre">pytorch_tabnet.metrics.Metric</span></code></a></p>
-<p>Unsupervised metric</p>
+<dl class="function">
+<dt id="pytorch_tabnet.multiclass_utils.type_of_target">
+<code class="sig-prename descclassname">pytorch_tabnet.multiclass_utils.</code><code class="sig-name descname">type_of_target</code><span class="sig-paren">(</span><em class="sig-param">y</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/multiclass_utils.html#type_of_target"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.multiclass_utils.type_of_target" title="Permalink to this definition">¶</a></dt>
+<dd><p>Determine the type of data indicated by the target.</p>
+<p>Note that this type is the most specific type that can be inferred.
+For example:</p>
+<blockquote>
+<div><ul class="simple">
+<li><p><code class="docutils literal notranslate"><span class="pre">binary</span></code> is more specific but compatible with <code class="docutils literal notranslate"><span class="pre">multiclass</span></code>.</p></li>
+<li><p><code class="docutils literal notranslate"><span class="pre">multiclass</span></code> of integers is more specific but compatible with
+<code class="docutils literal notranslate"><span class="pre">continuous</span></code>.</p></li>
+<li><p><code class="docutils literal notranslate"><span class="pre">multilabel-indicator</span></code> is more specific but compatible with
+<code class="docutils literal notranslate"><span class="pre">multiclass-multioutput</span></code>.</p></li>
+</ul>
+</div></blockquote>
+<dl class="field-list simple">
+<dt class="field-odd">Parameters</dt>
+<dd class="field-odd"><p><strong>y</strong> (<em>array-like</em>) – </p>
+</dd>
+<dt class="field-even">Returns</dt>
+<dd class="field-even"><p><p><strong>target_type</strong> – One of:</p>
+<ul class="simple">
+<li><p>’continuous’: <cite>y</cite> is an array-like of floats that are not all
+integers, and is 1d or a column vector.</p></li>
+<li><p>’continuous-multioutput’: <cite>y</cite> is a 2d array of floats that are
+not all integers, and both dimensions are of size &gt; 1.</p></li>
+<li><p>’binary’: <cite>y</cite> contains &lt;= 2 discrete values and is 1d or a column
+vector.</p></li>
+<li><p>’multiclass’: <cite>y</cite> contains more than two discrete values, is not a
+sequence of sequences, and is 1d or a column vector.</p></li>
+<li><p>’multiclass-multioutput’: <cite>y</cite> is a 2d array that contains more
+than two discrete values, is not a sequence of sequences, and both
+dimensions are of size &gt; 1.</p></li>
+<li><p>’multilabel-indicator’: <cite>y</cite> is a label indicator matrix, an array
+of two dimensions with at least two columns, and at most 2 unique
+values.</p></li>
+<li><p>’unknown’: <cite>y</cite> is array-like but none of the above, such as a 3d
+array, sequence of sequences, or an array of non-sequence objects.</p></li>
+</ul>
+</p>
+</dd>
+<dt class="field-odd">Return type</dt>
+<dd class="field-odd"><p>string</p>
+</dd>
+</dl>
+<p class="rubric">Examples</p>
+<div class="doctest highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="kn">import</span> <span class="nn">numpy</span> <span class="k">as</span> <span class="nn">np</span>
+<span class="gp">&gt;&gt;&gt; </span><span class="n">type_of_target</span><span class="p">([</span><span class="mf">0.1</span><span class="p">,</span> <span class="mf">0.6</span><span class="p">])</span>
+<span class="go">&#39;continuous&#39;</span>
+<span class="gp">&gt;&gt;&gt; </span><span class="n">type_of_target</span><span class="p">([</span><span class="mi">1</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="o">-</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">])</span>
+<span class="go">&#39;binary&#39;</span>
+<span class="gp">&gt;&gt;&gt; </span><span class="n">type_of_target</span><span class="p">([</span><span class="s1">&#39;a&#39;</span><span class="p">,</span> <span class="s1">&#39;b&#39;</span><span class="p">,</span> <span class="s1">&#39;a&#39;</span><span class="p">])</span>
+<span class="go">&#39;binary&#39;</span>
+<span class="gp">&gt;&gt;&gt; </span><span class="n">type_of_target</span><span class="p">([</span><span class="mf">1.0</span><span class="p">,</span> <span class="mf">2.0</span><span class="p">])</span>
+<span class="go">&#39;binary&#39;</span>
+<span class="gp">&gt;&gt;&gt; </span><span class="n">type_of_target</span><span class="p">([</span><span class="mi">1</span><span class="p">,</span> <span class="mi">0</span><span class="p">,</span> <span class="mi">2</span><span class="p">])</span>
+<span class="go">&#39;multiclass&#39;</span>
+<span class="gp">&gt;&gt;&gt; </span><span class="n">type_of_target</span><span class="p">([</span><span class="mf">1.0</span><span class="p">,</span> <span class="mf">0.0</span><span class="p">,</span> <span class="mf">3.0</span><span class="p">])</span>
+<span class="go">&#39;multiclass&#39;</span>
+<span class="gp">&gt;&gt;&gt; </span><span class="n">type_of_target</span><span class="p">([</span><span class="s1">&#39;a&#39;</span><span class="p">,</span> <span class="s1">&#39;b&#39;</span><span class="p">,</span> <span class="s1">&#39;c&#39;</span><span class="p">])</span>
+<span class="go">&#39;multiclass&#39;</span>
+<span class="gp">&gt;&gt;&gt; </span><span class="n">type_of_target</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">([[</span><span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">],</span> <span class="p">[</span><span class="mi">3</span><span class="p">,</span> <span class="mi">1</span><span class="p">]]))</span>
+<span class="go">&#39;multiclass-multioutput&#39;</span>
+<span class="gp">&gt;&gt;&gt; </span><span class="n">type_of_target</span><span class="p">([[</span><span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">]])</span>
+<span class="go">&#39;multiclass-multioutput&#39;</span>
+<span class="gp">&gt;&gt;&gt; </span><span class="n">type_of_target</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">([[</span><span class="mf">1.5</span><span class="p">,</span> <span class="mf">2.0</span><span class="p">],</span> <span class="p">[</span><span class="mf">3.0</span><span class="p">,</span> <span class="mf">1.6</span><span class="p">]]))</span>
+<span class="go">&#39;continuous-multioutput&#39;</span>
+<span class="gp">&gt;&gt;&gt; </span><span class="n">type_of_target</span><span class="p">(</span><span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">([[</span><span class="mi">0</span><span class="p">,</span> <span class="mi">1</span><span class="p">],</span> <span class="p">[</span><span class="mi">1</span><span class="p">,</span> <span class="mi">1</span><span class="p">]]))</span>
+<span class="go">&#39;multilabel-indicator&#39;</span>
+</pre></div>
+</div>
 </dd></dl>
 
 <dl class="function">
-<dt id="pytorch_tabnet.metrics.check_metrics">
-<code class="sig-prename descclassname">pytorch_tabnet.metrics.</code><code class="sig-name descname">check_metrics</code><span class="sig-paren">(</span><em class="sig-param">metrics</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/metrics.html#check_metrics"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.metrics.check_metrics" title="Permalink to this definition">¶</a></dt>
-<dd><p>Check if custom metrics are provided.</p>
+<dt id="pytorch_tabnet.multiclass_utils.unique_labels">
+<code class="sig-prename descclassname">pytorch_tabnet.multiclass_utils.</code><code class="sig-name descname">unique_labels</code><span class="sig-paren">(</span><em class="sig-param">*ys</em><span class="sig-paren">)</span><a class="reference internal" href="../_modules/pytorch_tabnet/multiclass_utils.html#unique_labels"><span class="viewcode-link">[source]</span></a><a class="headerlink" href="#pytorch_tabnet.multiclass_utils.unique_labels" title="Permalink to this definition">¶</a></dt>
+<dd><p>Extract an ordered array of unique labels</p>
+<dl class="simple">
+<dt>We don’t allow:</dt><dd><ul class="simple">
+<li><p>mix of multilabel and multiclass (single label) targets</p></li>
+<li><p>mix of label indicator matrix and anything else,
+because there are no explicit labels)</p></li>
+<li><p>mix of label indicator matrices of different sizes</p></li>
+<li><p>mix of string and integer labels</p></li>
+</ul>
+</dd>
+</dl>
+<p>At the moment, we also don’t allow “multiclass-multioutput” input type.</p>
 <dl class="field-list simple">
 <dt class="field-odd">Parameters</dt>
-<dd class="field-odd"><p><strong>metrics</strong> (<em>list of str</em><em> or </em><em>classes</em>) – List with built-in metrics (str) or custom metrics (classes).</p>
+<dd class="field-odd"><p><strong>*ys</strong> (<em>array-likes</em>) – </p>
 </dd>
 <dt class="field-even">Returns</dt>
-<dd class="field-even"><p><strong>val_metrics</strong> – List of metric names.</p>
+<dd class="field-even"><p><strong>out</strong> – An ordered array of unique labels.</p>
 </dd>
 <dt class="field-odd">Return type</dt>
-<dd class="field-odd"><p>list of str</p>
+<dd class="field-odd"><p>numpy array of shape [n_unique_labels]</p>
 </dd>
 </dl>
+<p class="rubric">Examples</p>
+<div class="doctest highlight-default notranslate"><div class="highlight"><pre><span></span><span class="gp">&gt;&gt;&gt; </span><span class="kn">from</span> <span class="nn">sklearn.utils.multiclass</span> <span class="kn">import</span> <span class="n">unique_labels</span>
+<span class="gp">&gt;&gt;&gt; </span><span class="n">unique_labels</span><span class="p">([</span><span class="mi">3</span><span class="p">,</span> <span class="mi">5</span><span class="p">,</span> <span class="mi">5</span><span class="p">,</span> <span class="mi">5</span><span class="p">,</span> <span class="mi">7</span><span class="p">,</span> <span class="mi">7</span><span class="p">])</span>
+<span class="go">array([3, 5, 7])</span>
+<span class="gp">&gt;&gt;&gt; </span><span class="n">unique_labels</span><span class="p">([</span><span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="mi">4</span><span class="p">],</span> <span class="p">[</span><span class="mi">2</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">3</span><span class="p">,</span> <span class="mi">4</span><span class="p">])</span>
+<span class="go">array([1, 2, 3, 4])</span>
+<span class="gp">&gt;&gt;&gt; </span><span class="n">unique_labels</span><span class="p">([</span><span class="mi">1</span><span class="p">,</span> <span class="mi">2</span><span class="p">,</span> <span class="mi">10</span><span class="p">],</span> <span class="p">[</span><span class="mi">5</span><span class="p">,</span> <span class="mi">11</span><span class="p">])</span>
+<span class="go">array([ 1,  2,  5, 10, 11])</span>
+</pre></div>
+</div>
 </dd></dl>
 
-</div>
-</div>
+</section>
+</section>
+</section>
 
 
            </div>
diff --git a/docs/genindex.html b/docs/genindex.html
index c8dd5366..3c9e444d 100644
--- a/docs/genindex.html
+++ b/docs/genindex.html
@@ -83,14 +83,17 @@
               
             
             
-              <p class="caption"><span class="caption-text">Contents:</span></p>
+              <p><span class="caption-text">Contents:</span></p>
 <ul>
 <li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html">README</a></li>
 <li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#tabnet-attentive-interpretable-tabular-learning">TabNet : Attentive Interpretable Tabular Learning</a></li>
 <li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#installation">Installation</a></li>
-<li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#what-problems-does-pytorch-tabnet-handles">What problems does pytorch-tabnet handles?</a></li>
+<li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#contributing">Contributing</a></li>
+<li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#what-problems-does-pytorch-tabnet-handle">What problems does pytorch-tabnet handle?</a></li>
 <li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#how-to-use-it">How to use it?</a></li>
 <li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#semi-supervised-pre-training">Semi-supervised pre-training</a></li>
+<li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#data-augmentation-on-the-fly">Data augmentation on the fly</a></li>
+<li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#easy-saving-and-loading">Easy saving and loading</a></li>
 <li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#useful-links">Useful links</a></li>
 <li class="toctree-l1"><a class="reference internal" href="generated_docs/pytorch_tabnet.html">pytorch_tabnet package</a></li>
 </ul>
@@ -253,15 +256,21 @@ <h2 id="C">C</h2>
         <li><a href="generated_docs/pytorch_tabnet.html#pytorch_tabnet.tab_model.TabNetRegressor.cat_idxs">(pytorch_tabnet.tab_model.TabNetRegressor attribute)</a>
 </li>
       </ul></li>
-  </ul></td>
-  <td style="width: 33%; vertical-align: top;"><ul>
       <li><a href="generated_docs/pytorch_tabnet.html#pytorch_tabnet.multiclass_utils.check_classification_targets">check_classification_targets() (in module pytorch_tabnet.multiclass_utils)</a>
 </li>
+      <li><a href="generated_docs/pytorch_tabnet.html#pytorch_tabnet.utils.check_input">check_input() (in module pytorch_tabnet.utils)</a>
+</li>
+  </ul></td>
+  <td style="width: 33%; vertical-align: top;"><ul>
       <li><a href="generated_docs/pytorch_tabnet.html#pytorch_tabnet.metrics.check_metrics">check_metrics() (in module pytorch_tabnet.metrics)</a>
 </li>
       <li><a href="generated_docs/pytorch_tabnet.html#pytorch_tabnet.multiclass_utils.check_output_dim">check_output_dim() (in module pytorch_tabnet.multiclass_utils)</a>
 </li>
       <li><a href="generated_docs/pytorch_tabnet.html#pytorch_tabnet.multiclass_utils.check_unique_type">check_unique_type() (in module pytorch_tabnet.multiclass_utils)</a>
+</li>
+      <li><a href="generated_docs/pytorch_tabnet.html#pytorch_tabnet.utils.check_warm_start">check_warm_start() (in module pytorch_tabnet.utils)</a>
+</li>
+      <li><a href="generated_docs/pytorch_tabnet.html#pytorch_tabnet.augmentations.ClassificationSMOTE">ClassificationSMOTE (class in pytorch_tabnet.augmentations)</a>
 </li>
       <li><a href="generated_docs/pytorch_tabnet.html#pytorch_tabnet.abstract_model.TabModel.clip_value">clip_value (pytorch_tabnet.abstract_model.TabModel attribute)</a>
 </li>
@@ -504,12 +513,16 @@ <h2 id="N">N</h2>
       <li><a href="generated_docs/pytorch_tabnet.html#pytorch_tabnet.abstract_model.TabModel.n_a">n_a (pytorch_tabnet.abstract_model.TabModel attribute)</a>
 </li>
       <li><a href="generated_docs/pytorch_tabnet.html#pytorch_tabnet.abstract_model.TabModel.n_d">n_d (pytorch_tabnet.abstract_model.TabModel attribute)</a>
+</li>
+      <li><a href="generated_docs/pytorch_tabnet.html#pytorch_tabnet.abstract_model.TabModel.n_indep_decoder">n_indep_decoder (pytorch_tabnet.abstract_model.TabModel attribute)</a>
 </li>
   </ul></td>
   <td style="width: 33%; vertical-align: top;"><ul>
       <li><a href="generated_docs/pytorch_tabnet.html#pytorch_tabnet.abstract_model.TabModel.n_independent">n_independent (pytorch_tabnet.abstract_model.TabModel attribute)</a>
 </li>
       <li><a href="generated_docs/pytorch_tabnet.html#pytorch_tabnet.abstract_model.TabModel.n_shared">n_shared (pytorch_tabnet.abstract_model.TabModel attribute)</a>
+</li>
+      <li><a href="generated_docs/pytorch_tabnet.html#pytorch_tabnet.abstract_model.TabModel.n_shared_decoder">n_shared_decoder (pytorch_tabnet.abstract_model.TabModel attribute)</a>
 </li>
       <li><a href="generated_docs/pytorch_tabnet.html#pytorch_tabnet.abstract_model.TabModel.n_steps">n_steps (pytorch_tabnet.abstract_model.TabModel attribute)</a>
 </li>
@@ -642,6 +655,8 @@ <h2 id="P">P</h2>
   </ul></td>
   <td style="width: 33%; vertical-align: top;"><ul>
       <li><a href="generated_docs/pytorch_tabnet.html#module-pytorch_tabnet.abstract_model">pytorch_tabnet.abstract_model (module)</a>
+</li>
+      <li><a href="generated_docs/pytorch_tabnet.html#module-pytorch_tabnet.augmentations">pytorch_tabnet.augmentations (module)</a>
 </li>
       <li><a href="generated_docs/pytorch_tabnet.html#module-pytorch_tabnet.callbacks">pytorch_tabnet.callbacks (module)</a>
 </li>
@@ -670,6 +685,8 @@ <h2 id="R">R</h2>
 <table style="width: 100%" class="indextable genindextable"><tr>
   <td style="width: 33%; vertical-align: top;"><ul>
       <li><a href="generated_docs/pytorch_tabnet.html#pytorch_tabnet.tab_network.RandomObfuscator">RandomObfuscator (class in pytorch_tabnet.tab_network)</a>
+</li>
+      <li><a href="generated_docs/pytorch_tabnet.html#pytorch_tabnet.augmentations.RegressionSMOTE">RegressionSMOTE (class in pytorch_tabnet.augmentations)</a>
 </li>
   </ul></td>
   <td style="width: 33%; vertical-align: top;"><ul>
@@ -782,13 +799,17 @@ <h2 id="U">U</h2>
       <li><a href="generated_docs/pytorch_tabnet.html#pytorch_tabnet.multiclass_utils.unique_labels">unique_labels() (in module pytorch_tabnet.multiclass_utils)</a>
 </li>
       <li><a href="generated_docs/pytorch_tabnet.html#pytorch_tabnet.metrics.UnsupervisedLoss">UnsupervisedLoss() (in module pytorch_tabnet.metrics)</a>
+</li>
+      <li><a href="generated_docs/pytorch_tabnet.html#pytorch_tabnet.metrics.UnsupervisedLossNumpy">UnsupervisedLossNumpy() (in module pytorch_tabnet.metrics)</a>
 </li>
       <li><a href="generated_docs/pytorch_tabnet.html#pytorch_tabnet.metrics.UnsupervisedMetric">UnsupervisedMetric (class in pytorch_tabnet.metrics)</a>
 </li>
-      <li><a href="generated_docs/pytorch_tabnet.html#pytorch_tabnet.metrics.UnsupMetricContainer">UnsupMetricContainer (class in pytorch_tabnet.metrics)</a>
+      <li><a href="generated_docs/pytorch_tabnet.html#pytorch_tabnet.metrics.UnsupervisedNumpyMetric">UnsupervisedNumpyMetric (class in pytorch_tabnet.metrics)</a>
 </li>
   </ul></td>
   <td style="width: 33%; vertical-align: top;"><ul>
+      <li><a href="generated_docs/pytorch_tabnet.html#pytorch_tabnet.metrics.UnsupMetricContainer">UnsupMetricContainer (class in pytorch_tabnet.metrics)</a>
+</li>
       <li><a href="generated_docs/pytorch_tabnet.html#pytorch_tabnet.abstract_model.TabModel.update_fit_params">update_fit_params() (pytorch_tabnet.abstract_model.TabModel method)</a>
 
       <ul>
diff --git a/docs/index.html b/docs/index.html
index 454f6f35..27c3e93a 100644
--- a/docs/index.html
+++ b/docs/index.html
@@ -4,7 +4,8 @@
 <html class="writer-html5" lang="en" >
 <head>
   <meta charset="utf-8">
-  
+  <meta name="generator" content="Docutils 0.17.1: http://docutils.sourceforge.net/" />
+
   <meta name="viewport" content="width=device-width, initial-scale=1.0">
   
   <title>Welcome to pytorch_tabnet’s documentation! &mdash; pytorch_tabnet  documentation</title>
@@ -83,14 +84,17 @@
               
             
             
-              <p class="caption"><span class="caption-text">Contents:</span></p>
+              <p><span class="caption-text">Contents:</span></p>
 <ul>
 <li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html">README</a></li>
 <li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#tabnet-attentive-interpretable-tabular-learning">TabNet : Attentive Interpretable Tabular Learning</a></li>
 <li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#installation">Installation</a></li>
-<li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#what-problems-does-pytorch-tabnet-handles">What problems does pytorch-tabnet handles?</a></li>
+<li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#contributing">Contributing</a></li>
+<li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#what-problems-does-pytorch-tabnet-handle">What problems does pytorch-tabnet handle?</a></li>
 <li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#how-to-use-it">How to use it?</a></li>
 <li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#semi-supervised-pre-training">Semi-supervised pre-training</a></li>
+<li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#data-augmentation-on-the-fly">Data augmentation on the fly</a></li>
+<li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#easy-saving-and-loading">Easy saving and loading</a></li>
 <li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#useful-links">Useful links</a></li>
 <li class="toctree-l1"><a class="reference internal" href="generated_docs/pytorch_tabnet.html">pytorch_tabnet package</a></li>
 </ul>
@@ -158,10 +162,10 @@
           <div role="main" class="document" itemscope="itemscope" itemtype="http://schema.org/Article">
            <div itemprop="articleBody">
             
-  <div class="section" id="welcome-to-pytorch-tabnet-s-documentation">
+  <section id="welcome-to-pytorch-tabnet-s-documentation">
 <h1>Welcome to pytorch_tabnet’s documentation!<a class="headerlink" href="#welcome-to-pytorch-tabnet-s-documentation" title="Permalink to this headline">¶</a></h1>
 <div class="toctree-wrapper compound">
-<p class="caption"><span class="caption-text">Contents:</span></p>
+<p><span class="caption-text">Contents:</span></p>
 <ul>
 <li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html">README</a></li>
 <li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#tabnet-attentive-interpretable-tabular-learning">TabNet : Attentive Interpretable Tabular Learning</a></li>
@@ -170,43 +174,47 @@ <h1>Welcome to pytorch_tabnet’s documentation!<a class="headerlink" href="#wel
 <li class="toctree-l2"><a class="reference internal" href="generated_docs/README.html#source-code">Source code</a></li>
 </ul>
 </li>
-<li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#what-problems-does-pytorch-tabnet-handles">What problems does pytorch-tabnet handles?</a></li>
+<li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#contributing">Contributing</a></li>
+<li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#what-problems-does-pytorch-tabnet-handle">What problems does pytorch-tabnet handle?</a></li>
 <li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#how-to-use-it">How to use it?</a><ul>
 <li class="toctree-l2"><a class="reference internal" href="generated_docs/README.html#default-eval-metric">Default eval_metric</a></li>
 <li class="toctree-l2"><a class="reference internal" href="generated_docs/README.html#custom-evaluation-metrics">Custom evaluation metrics</a></li>
 </ul>
 </li>
 <li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#semi-supervised-pre-training">Semi-supervised pre-training</a></li>
+<li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#data-augmentation-on-the-fly">Data augmentation on the fly</a></li>
+<li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#easy-saving-and-loading">Easy saving and loading</a></li>
 <li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#useful-links">Useful links</a><ul>
 <li class="toctree-l2"><a class="reference internal" href="generated_docs/README.html#model-parameters">Model parameters</a></li>
 <li class="toctree-l2"><a class="reference internal" href="generated_docs/README.html#fit-parameters">Fit parameters</a></li>
 </ul>
 </li>
 <li class="toctree-l1"><a class="reference internal" href="generated_docs/pytorch_tabnet.html">pytorch_tabnet package</a><ul>
-<li class="toctree-l2"><a class="reference internal" href="generated_docs/pytorch_tabnet.html#module-pytorch_tabnet.utils">pytorch_tabnet.utils module</a></li>
 <li class="toctree-l2"><a class="reference internal" href="generated_docs/pytorch_tabnet.html#module-pytorch_tabnet.pretraining_utils">pytorch_tabnet.pretraining_utils module</a></li>
+<li class="toctree-l2"><a class="reference internal" href="generated_docs/pytorch_tabnet.html#module-pytorch_tabnet.augmentations">pytorch_tabnet.augmentations module</a></li>
 <li class="toctree-l2"><a class="reference internal" href="generated_docs/pytorch_tabnet.html#module-pytorch_tabnet.tab_network">pytorch_tabnet.tab_network module</a></li>
-<li class="toctree-l2"><a class="reference internal" href="generated_docs/pytorch_tabnet.html#module-pytorch_tabnet.multitask">pytorch_tabnet.multitask module</a></li>
-<li class="toctree-l2"><a class="reference internal" href="generated_docs/pytorch_tabnet.html#module-pytorch_tabnet.multiclass_utils">pytorch_tabnet.multiclass_utils module</a></li>
+<li class="toctree-l2"><a class="reference internal" href="generated_docs/pytorch_tabnet.html#module-pytorch_tabnet.metrics">pytorch_tabnet.metrics module</a></li>
 <li class="toctree-l2"><a class="reference internal" href="generated_docs/pytorch_tabnet.html#module-pytorch_tabnet.tab_model">pytorch_tabnet.tab_model module</a></li>
-<li class="toctree-l2"><a class="reference internal" href="generated_docs/pytorch_tabnet.html#module-pytorch_tabnet.abstract_model">pytorch_tabnet.abstract_model module</a></li>
+<li class="toctree-l2"><a class="reference internal" href="generated_docs/pytorch_tabnet.html#module-pytorch_tabnet.sparsemax">pytorch_tabnet.sparsemax module</a></li>
 <li class="toctree-l2"><a class="reference internal" href="generated_docs/pytorch_tabnet.html#module-pytorch_tabnet.callbacks">pytorch_tabnet.callbacks module</a></li>
+<li class="toctree-l2"><a class="reference internal" href="generated_docs/pytorch_tabnet.html#module-pytorch_tabnet.abstract_model">pytorch_tabnet.abstract_model module</a></li>
 <li class="toctree-l2"><a class="reference internal" href="generated_docs/pytorch_tabnet.html#module-pytorch_tabnet.pretraining">pytorch_tabnet.pretraining module</a></li>
-<li class="toctree-l2"><a class="reference internal" href="generated_docs/pytorch_tabnet.html#module-pytorch_tabnet.sparsemax">pytorch_tabnet.sparsemax module</a></li>
-<li class="toctree-l2"><a class="reference internal" href="generated_docs/pytorch_tabnet.html#module-pytorch_tabnet.metrics">pytorch_tabnet.metrics module</a></li>
+<li class="toctree-l2"><a class="reference internal" href="generated_docs/pytorch_tabnet.html#module-pytorch_tabnet.utils">pytorch_tabnet.utils module</a></li>
+<li class="toctree-l2"><a class="reference internal" href="generated_docs/pytorch_tabnet.html#module-pytorch_tabnet.multitask">pytorch_tabnet.multitask module</a></li>
+<li class="toctree-l2"><a class="reference internal" href="generated_docs/pytorch_tabnet.html#module-pytorch_tabnet.multiclass_utils">pytorch_tabnet.multiclass_utils module</a></li>
 </ul>
 </li>
 </ul>
 </div>
-</div>
-<div class="section" id="indices-and-tables">
+</section>
+<section id="indices-and-tables">
 <h1>Indices and tables<a class="headerlink" href="#indices-and-tables" title="Permalink to this headline">¶</a></h1>
 <ul class="simple">
 <li><p><a class="reference internal" href="genindex.html"><span class="std std-ref">Index</span></a></p></li>
 <li><p><a class="reference internal" href="py-modindex.html"><span class="std std-ref">Module Index</span></a></p></li>
 <li><p><a class="reference internal" href="search.html"><span class="std std-ref">Search Page</span></a></p></li>
 </ul>
-</div>
+</section>
 
 
            </div>
diff --git a/docs/py-modindex.html b/docs/py-modindex.html
index 33811ddb..6dee3c89 100644
--- a/docs/py-modindex.html
+++ b/docs/py-modindex.html
@@ -85,14 +85,17 @@
               
             
             
-              <p class="caption"><span class="caption-text">Contents:</span></p>
+              <p><span class="caption-text">Contents:</span></p>
 <ul>
 <li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html">README</a></li>
 <li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#tabnet-attentive-interpretable-tabular-learning">TabNet : Attentive Interpretable Tabular Learning</a></li>
 <li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#installation">Installation</a></li>
-<li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#what-problems-does-pytorch-tabnet-handles">What problems does pytorch-tabnet handles?</a></li>
+<li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#contributing">Contributing</a></li>
+<li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#what-problems-does-pytorch-tabnet-handle">What problems does pytorch-tabnet handle?</a></li>
 <li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#how-to-use-it">How to use it?</a></li>
 <li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#semi-supervised-pre-training">Semi-supervised pre-training</a></li>
+<li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#data-augmentation-on-the-fly">Data augmentation on the fly</a></li>
+<li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#easy-saving-and-loading">Easy saving and loading</a></li>
 <li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#useful-links">Useful links</a></li>
 <li class="toctree-l1"><a class="reference internal" href="generated_docs/pytorch_tabnet.html">pytorch_tabnet package</a></li>
 </ul>
@@ -178,6 +181,11 @@ <h1>Python Module Index</h1>
        <td>&#160;&#160;&#160;
        <a href="generated_docs/pytorch_tabnet.html#module-pytorch_tabnet.abstract_model"><code class="xref">pytorch_tabnet.abstract_model</code></a></td><td>
        <em></em></td></tr>
+     <tr class="cg-1">
+       <td></td>
+       <td>&#160;&#160;&#160;
+       <a href="generated_docs/pytorch_tabnet.html#module-pytorch_tabnet.augmentations"><code class="xref">pytorch_tabnet.augmentations</code></a></td><td>
+       <em></em></td></tr>
      <tr class="cg-1">
        <td></td>
        <td>&#160;&#160;&#160;
diff --git a/docs/search.html b/docs/search.html
index ab0e1e5f..54ba30a5 100644
--- a/docs/search.html
+++ b/docs/search.html
@@ -84,14 +84,17 @@
               
             
             
-              <p class="caption"><span class="caption-text">Contents:</span></p>
+              <p><span class="caption-text">Contents:</span></p>
 <ul>
 <li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html">README</a></li>
 <li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#tabnet-attentive-interpretable-tabular-learning">TabNet : Attentive Interpretable Tabular Learning</a></li>
 <li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#installation">Installation</a></li>
-<li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#what-problems-does-pytorch-tabnet-handles">What problems does pytorch-tabnet handles?</a></li>
+<li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#contributing">Contributing</a></li>
+<li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#what-problems-does-pytorch-tabnet-handle">What problems does pytorch-tabnet handle?</a></li>
 <li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#how-to-use-it">How to use it?</a></li>
 <li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#semi-supervised-pre-training">Semi-supervised pre-training</a></li>
+<li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#data-augmentation-on-the-fly">Data augmentation on the fly</a></li>
+<li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#easy-saving-and-loading">Easy saving and loading</a></li>
 <li class="toctree-l1"><a class="reference internal" href="generated_docs/README.html#useful-links">Useful links</a></li>
 <li class="toctree-l1"><a class="reference internal" href="generated_docs/pytorch_tabnet.html">pytorch_tabnet package</a></li>
 </ul>
diff --git a/docs/searchindex.js b/docs/searchindex.js
index 4d732e3f..b069376c 100644
--- a/docs/searchindex.js
+++ b/docs/searchindex.js
@@ -1 +1 @@
-Search.setIndex({docnames:["generated_docs/README","generated_docs/pytorch_tabnet","index"],envversion:{"sphinx.domains.c":1,"sphinx.domains.changeset":1,"sphinx.domains.citation":1,"sphinx.domains.cpp":1,"sphinx.domains.index":1,"sphinx.domains.javascript":1,"sphinx.domains.math":2,"sphinx.domains.python":1,"sphinx.domains.rst":1,"sphinx.domains.std":1,"sphinx.ext.viewcode":1,sphinx:56},filenames:["generated_docs/README.md","generated_docs/pytorch_tabnet.rst","index.rst"],objects:{"pytorch_tabnet.abstract_model":{TabModel:[1,1,1,""]},"pytorch_tabnet.abstract_model.TabModel":{cat_dims:[1,2,1,""],cat_emb_dim:[1,2,1,""],cat_idxs:[1,2,1,""],clip_value:[1,2,1,""],compute_loss:[1,3,1,""],device_name:[1,2,1,""],epsilon:[1,2,1,""],explain:[1,3,1,""],fit:[1,3,1,""],gamma:[1,2,1,""],input_dim:[1,2,1,""],lambda_sparse:[1,2,1,""],load_class_attrs:[1,3,1,""],load_model:[1,3,1,""],load_weights_from_unsupervised:[1,3,1,""],mask_type:[1,2,1,""],momentum:[1,2,1,""],n_a:[1,2,1,""],n_d:[1,2,1,""],n_independent:[1,2,1,""],n_shared:[1,2,1,""],n_steps:[1,2,1,""],optimizer_fn:[1,2,1,""],optimizer_params:[1,2,1,""],output_dim:[1,2,1,""],predict:[1,3,1,""],prepare_target:[1,3,1,""],save_model:[1,3,1,""],scheduler_fn:[1,2,1,""],scheduler_params:[1,2,1,""],seed:[1,2,1,""],update_fit_params:[1,3,1,""],verbose:[1,2,1,""]},"pytorch_tabnet.callbacks":{Callback:[1,1,1,""],CallbackContainer:[1,1,1,""],EarlyStopping:[1,1,1,""],History:[1,1,1,""],LRSchedulerCallback:[1,1,1,""]},"pytorch_tabnet.callbacks.Callback":{on_batch_begin:[1,3,1,""],on_batch_end:[1,3,1,""],on_epoch_begin:[1,3,1,""],on_epoch_end:[1,3,1,""],on_train_begin:[1,3,1,""],on_train_end:[1,3,1,""],set_params:[1,3,1,""],set_trainer:[1,3,1,""]},"pytorch_tabnet.callbacks.CallbackContainer":{append:[1,3,1,""],callbacks:[1,2,1,""],on_batch_begin:[1,3,1,""],on_batch_end:[1,3,1,""],on_epoch_begin:[1,3,1,""],on_epoch_end:[1,3,1,""],on_train_begin:[1,3,1,""],on_train_end:[1,3,1,""],set_params:[1,3,1,""],set_trainer:[1,3,1,""]},"pytorch_tabnet.callbacks.EarlyStopping":{early_stopping_metric:[1,2,1,""],is_maximize:[1,2,1,""],on_epoch_end:[1,3,1,""],on_train_end:[1,3,1,""],patience:[1,2,1,""],tol:[1,2,1,""]},"pytorch_tabnet.callbacks.History":{on_batch_end:[1,3,1,""],on_epoch_begin:[1,3,1,""],on_epoch_end:[1,3,1,""],on_train_begin:[1,3,1,""],trainer:[1,2,1,""],verbose:[1,2,1,""]},"pytorch_tabnet.callbacks.LRSchedulerCallback":{early_stopping_metric:[1,2,1,""],is_batch_level:[1,2,1,""],on_batch_end:[1,3,1,""],on_epoch_end:[1,3,1,""],optimizer:[1,2,1,""],scheduler_fn:[1,2,1,""],scheduler_params:[1,2,1,""]},"pytorch_tabnet.metrics":{AUC:[1,1,1,""],Accuracy:[1,1,1,""],BalancedAccuracy:[1,1,1,""],LogLoss:[1,1,1,""],MAE:[1,1,1,""],MSE:[1,1,1,""],Metric:[1,1,1,""],MetricContainer:[1,1,1,""],RMSE:[1,1,1,""],RMSLE:[1,1,1,""],UnsupMetricContainer:[1,1,1,""],UnsupervisedLoss:[1,4,1,""],UnsupervisedMetric:[1,1,1,""],check_metrics:[1,4,1,""]},"pytorch_tabnet.metrics.Metric":{get_metrics_by_names:[1,3,1,""]},"pytorch_tabnet.metrics.MetricContainer":{metric_names:[1,2,1,""],prefix:[1,2,1,""]},"pytorch_tabnet.metrics.UnsupMetricContainer":{metric_names:[1,2,1,""],prefix:[1,2,1,""]},"pytorch_tabnet.multiclass_utils":{assert_all_finite:[1,4,1,""],check_classification_targets:[1,4,1,""],check_output_dim:[1,4,1,""],check_unique_type:[1,4,1,""],infer_multitask_output:[1,4,1,""],infer_output_dim:[1,4,1,""],is_multilabel:[1,4,1,""],type_of_target:[1,4,1,""],unique_labels:[1,4,1,""]},"pytorch_tabnet.multitask":{TabNetMultiTaskClassifier:[1,1,1,""]},"pytorch_tabnet.multitask.TabNetMultiTaskClassifier":{cat_dims:[1,2,1,""],cat_idxs:[1,2,1,""],compute_loss:[1,3,1,""],optimizer_params:[1,2,1,""],predict:[1,3,1,""],predict_proba:[1,3,1,""],prepare_target:[1,3,1,""],scheduler_params:[1,2,1,""],stack_batches:[1,3,1,""],update_fit_params:[1,3,1,""]},"pytorch_tabnet.pretraining":{TabNetPretrainer:[1,1,1,""]},"pytorch_tabnet.pretraining.TabNetPretrainer":{cat_dims:[1,2,1,""],cat_idxs:[1,2,1,""],compute_loss:[1,3,1,""],fit:[1,3,1,""],optimizer_params:[1,2,1,""],predict:[1,3,1,""],prepare_target:[1,3,1,""],scheduler_params:[1,2,1,""],stack_batches:[1,3,1,""],update_fit_params:[1,3,1,""]},"pytorch_tabnet.pretraining_utils":{create_dataloaders:[1,4,1,""],validate_eval_set:[1,4,1,""]},"pytorch_tabnet.sparsemax":{Entmax15:[1,1,1,""],Entmax15Function:[1,1,1,""],Entmoid15:[1,1,1,""],Sparsemax:[1,1,1,""],SparsemaxFunction:[1,1,1,""],entmax15:[1,4,1,""],entmoid15:[1,4,1,""],sparsemax:[1,4,1,""]},"pytorch_tabnet.sparsemax.Entmax15":{forward:[1,3,1,""]},"pytorch_tabnet.sparsemax.Entmax15Function":{backward:[1,3,1,""],forward:[1,3,1,""]},"pytorch_tabnet.sparsemax.Entmoid15":{backward:[1,3,1,""],forward:[1,3,1,""]},"pytorch_tabnet.sparsemax.Sparsemax":{forward:[1,3,1,""]},"pytorch_tabnet.sparsemax.SparsemaxFunction":{backward:[1,3,1,""],forward:[1,3,1,""]},"pytorch_tabnet.tab_model":{TabNetClassifier:[1,1,1,""],TabNetRegressor:[1,1,1,""]},"pytorch_tabnet.tab_model.TabNetClassifier":{cat_dims:[1,2,1,""],cat_idxs:[1,2,1,""],compute_loss:[1,3,1,""],optimizer_params:[1,2,1,""],predict_func:[1,3,1,""],predict_proba:[1,3,1,""],prepare_target:[1,3,1,""],scheduler_params:[1,2,1,""],stack_batches:[1,3,1,""],update_fit_params:[1,3,1,""],weight_updater:[1,3,1,""]},"pytorch_tabnet.tab_model.TabNetRegressor":{cat_dims:[1,2,1,""],cat_idxs:[1,2,1,""],compute_loss:[1,3,1,""],optimizer_params:[1,2,1,""],predict_func:[1,3,1,""],prepare_target:[1,3,1,""],scheduler_params:[1,2,1,""],stack_batches:[1,3,1,""],update_fit_params:[1,3,1,""]},"pytorch_tabnet.tab_network":{AttentiveTransformer:[1,1,1,""],EmbeddingGenerator:[1,1,1,""],FeatTransformer:[1,1,1,""],GBN:[1,1,1,""],GLU_Block:[1,1,1,""],GLU_Layer:[1,1,1,""],RandomObfuscator:[1,1,1,""],TabNet:[1,1,1,""],TabNetDecoder:[1,1,1,""],TabNetEncoder:[1,1,1,""],TabNetNoEmbeddings:[1,1,1,""],TabNetPretraining:[1,1,1,""],initialize_glu:[1,4,1,""],initialize_non_glu:[1,4,1,""]},"pytorch_tabnet.tab_network.AttentiveTransformer":{forward:[1,3,1,""]},"pytorch_tabnet.tab_network.EmbeddingGenerator":{forward:[1,3,1,""]},"pytorch_tabnet.tab_network.FeatTransformer":{forward:[1,3,1,""]},"pytorch_tabnet.tab_network.GBN":{forward:[1,3,1,""]},"pytorch_tabnet.tab_network.GLU_Block":{forward:[1,3,1,""]},"pytorch_tabnet.tab_network.GLU_Layer":{forward:[1,3,1,""]},"pytorch_tabnet.tab_network.RandomObfuscator":{forward:[1,3,1,""]},"pytorch_tabnet.tab_network.TabNet":{forward:[1,3,1,""],forward_masks:[1,3,1,""]},"pytorch_tabnet.tab_network.TabNetDecoder":{forward:[1,3,1,""]},"pytorch_tabnet.tab_network.TabNetEncoder":{forward:[1,3,1,""],forward_masks:[1,3,1,""]},"pytorch_tabnet.tab_network.TabNetNoEmbeddings":{forward:[1,3,1,""],forward_masks:[1,3,1,""]},"pytorch_tabnet.tab_network.TabNetPretraining":{forward:[1,3,1,""],forward_masks:[1,3,1,""]},"pytorch_tabnet.utils":{ComplexEncoder:[1,1,1,""],PredictDataset:[1,1,1,""],TorchDataset:[1,1,1,""],create_dataloaders:[1,4,1,""],create_explain_matrix:[1,4,1,""],create_sampler:[1,4,1,""],define_device:[1,4,1,""],filter_weights:[1,4,1,""],validate_eval_set:[1,4,1,""]},"pytorch_tabnet.utils.ComplexEncoder":{"default":[1,3,1,""]},pytorch_tabnet:{abstract_model:[1,0,0,"-"],callbacks:[1,0,0,"-"],metrics:[1,0,0,"-"],multiclass_utils:[1,0,0,"-"],multitask:[1,0,0,"-"],pretraining:[1,0,0,"-"],pretraining_utils:[1,0,0,"-"],sparsemax:[1,0,0,"-"],tab_model:[1,0,0,"-"],tab_network:[1,0,0,"-"],utils:[1,0,0,"-"]}},objnames:{"0":["py","module","Python module"],"1":["py","class","Python class"],"2":["py","attribute","Python attribute"],"3":["py","method","Python method"],"4":["py","function","Python function"]},objtypes:{"0":"py:module","1":"py:class","2":"py:attribute","3":"py:method","4":"py:function"},terms:{"1st":0,"abstract":1,"boolean":1,"case":1,"class":0,"default":[1,2],"float":[0,1],"function":0,"import":[0,1],"int":[0,1],"new":[0,1],"return":[0,1],"static":1,"throw":1,"true":[0,1],"try":[0,1],"while":1,Added:0,For:1,One:1,The:[0,1],Use:1,Useful:2,Using:1,__call__:0,__init__:0,_contextmethodmixin:1,_maxim:0,_name:0,a_max:0,a_min:0,abov:1,abs:1,absolut:1,abstract_model:2,accept:1,accord:[0,1],accuraci:[0,1],adam:[0,1],after:1,afterward:1,alia:1,all:[0,1],allow:1,allow_nan:1,along:1,alpha:1,also:1,although:1,amount:1,ani:[0,1],anyth:1,append:1,appli:[0,1],arbitrari:1,architectur:0,argument:1,arik:0,arrai:[0,1],arxiv:[0,1],assert_all_finit:1,assign:0,astudillo:1,attent:[1,2],attentivetransform:1,attribut:1,auc:[0,1],auto:[0,1],autograd:1,autom:[0,1],automat:[0,1],avail:0,averag:1,avoid:1,backward:1,balanc:[0,1],balancedaccuraci:1,base:1,baseestim:1,batch:[0,1],batch_out:[],batch_siz:[0,1],becaus:[0,1],been:0,befor:[0,1],bellow:0,ben:1,best:0,better:0,between:[0,1],bigger:0,binari:[0,1],blob:0,block:1,bool:[0,1],both:1,build:[0,1],built:1,call:1,callabl:1,callback:[0,2],callbackcontain:1,can:[0,1],capac:0,care:1,cat:[],cat_dim:[0,1],cat_emb_dim:[0,1],cat_idx:[0,1],categor:[0,1],censu:[],certain:1,chang:[0,1],check:1,check_circular:1,check_classification_target:1,check_metr:1,check_nan:[],check_output_dim:1,check_unique_typ:1,choic:0,cite:1,class_attr:1,classic:[0,1],classif:[0,1],classmethod:1,clf:0,clip:[0,1],clip_valu:[0,1],clone:0,close:0,cls:1,code:2,coeffici:0,column:1,com:[0,1],compat:[0,1],complet:0,complexencod:1,comput:1,compute_loss:1,consecut:[0,1],contain:[0,1],content:2,context:1,continu:1,contribut:0,convert:1,corr:1,correct:1,correl:0,correspond:1,could:[0,1],counter:1,cpu:1,creat:[0,1],create_dataload:1,create_explain_matrix:1,create_sampl:1,cross:0,ctx:1,cuda:1,current:0,custom:[1,2],data:[0,1],dataload:[0,1],dataset:1,dblp:1,decai:0,decis:0,deduc:0,deep:1,deeprecomodel:1,def:[0,1],defin:[0,1],define_devic:1,degener:1,depend:[0,1],descript:1,detail:1,detect:[0,1],determin:1,develop:0,devic:1,device_nam:[0,1],dict:[0,1],dictionari:1,dictionnari:[0,1],did:1,differ:1,differenti:1,difficulti:0,dim:1,dimens:1,discret:1,disk:0,distinct:1,divid:0,divis:1,docker:0,doe:[1,2],doing:0,don:1,dreamquark:0,dreamquarktabnet:0,drop:[0,1],drop_last:[0,1],dure:[0,1],each:[0,1],earli:[0,1],early_stopping_metr:1,earlystop:1,easi:2,easier:0,easili:0,either:[0,1],element:1,els:1,emb:[],embded:[],embed:[0,1],embedded_x:1,embeddinggener:1,enabl:0,encod:1,end:0,ensur:1,ensure_ascii:1,entmax15:1,entmax15funct:1,entmax:[0,1],entmoid15:1,entropi:0,epoch:[0,1],eps:1,epsilon:[0,1],equal:1,equival:1,error:1,eval:[0,1,2],eval_metr:1,eval_nam:[0,1],eval_set:[0,1],evalu:[1,2],event:1,everi:[0,1],exact:1,exampl:[0,1],except:1,exit:1,expert:1,explain:1,explan:1,explanatori:0,explicit:1,extra:0,extract:1,factori:1,fals:[0,1],feattransform:1,featu:[],featur:[0,1],few:0,file:1,filepath:1,filter_weight:1,first:1,fit:[1,2],follow:[0,1],forest:[],format:1,former:1,formula:1,forward:1,forward_mask:1,found:0,frequenc:1,from:[0,1],from_unsupervis:[0,1],gamma:[0,1],gate:0,gbn:1,gener:1,get:[0,1],get_metrics_by_nam:1,ghost:[0,1],gini:0,git:0,github:[0,1],give:0,given:[0,1],glu:1,glu_block:1,glu_lay:1,good:0,gpu:1,grad_output:1,gradient:[0,1],greater:0,handl:2,happen:1,harder:0,has:[0,1],have:[0,1],help:0,here:0,highli:1,histori:1,hold:1,hook:1,hot:1,how:[1,2],html:1,http:[0,1],idx:[],ignor:1,imeplement:[],implement:[0,1],improv:[0,1],includ:0,incomplet:1,indent:1,independ:[0,1],index:[1,2],indic:[0,1],infer:1,infer_multitask_output:1,infer_output_dim:1,infin:1,initi:[0,1],initialize_glu:1,initialize_non_glu:1,input:[0,1],input_dim:1,insid:0,instal:2,instanc:1,instead:1,integ:[0,1],interpret:2,invers:[0,1],ipynb:0,is_batch_level:1,is_maxim:1,is_multilabel:1,iter:1,its:0,join:0,journal:1,json:1,jsonencod:1,jupyt:0,kaggl:0,kei:[0,1],labda:[],lambda:1,lambda_spars:[0,1],larg:0,last:[0,1],later:0,latter:1,layer:0,learn:[1,2],least:[0,1],left:[0,1],length:[0,1],let:1,like:1,line:0,linear:0,link:2,list:[0,1],list_embedded_x:1,list_obfusc:1,list_output:1,list_y_scor:1,list_y_tru:1,load:[0,1],load_class_attr:1,load_model:1,load_weights_from_unsupervis:1,local:[0,1],log:1,logarithm:1,logloss:[0,1],longtensor:1,loop:1,loss:[0,1],loss_fn:[0,1],lower:0,lr_schedul:[0,1],lrschedulercallback:1,m_explain:1,mae:[0,1],main:1,make:[0,1],mandatori:0,mani:1,manual:1,map:1,martin:1,martinsa16:1,mask:[0,1],mask_typ:[0,1],match:0,matric:1,matrix:1,max:0,max_epoch:[0,1],maxim:[0,1],maximum:[0,1],mean:[0,1],mean_squared_log_error:1,memori:1,mention:0,method:1,metric:2,metric_nam:1,metriccontain:1,might:0,mini:0,minimum:1,mix:1,moa:0,modal:0,model:[1,2],model_nam:0,modul:2,moment:1,momentum:[0,1],monitor:1,more:[0,1],most:1,mse:[0,1],multi:0,multiclass:[0,1],multiclass_util:2,multilabel:1,multioutput:1,multipl:1,multitask:[0,2],must:[0,1],n_a:[0,1],n_d:[0,1],n_glu:1,n_glu_independ:1,n_independ:[0,1],n_sampl:1,n_share:[0,1],n_step:[0,1],n_unique_label:1,name:[0,1],nan:1,ndarrai:1,need:[0,1],needs_input_grad:1,neg:[0,1],network:1,neural:1,nicula:1,non:1,none:[0,1],normal:[0,1],note:[0,1],notebook:0,now:0,num:[],num_work:[0,1],number:[0,1],numpi:1,obf_var:1,obfusc:1,obj:1,object:1,occur:[],occurr:0,on_batch_begin:1,on_batch_end:1,on_epoch_begin:1,on_epoch_end:1,on_train_begin:1,on_train_end:1,one:[0,1],onecyclelr:1,ones:0,onli:1,oper:1,optim:[0,1],optimizer_fn:[0,1],optimizer_param:[0,1],optimo:[],option:0,order:[0,1],org:[0,1],orgin:1,origin:[0,1],other:1,otherwis:1,our:0,out:1,output:1,output_dim:1,over:1,overfit:0,overridden:1,overwritten:0,own:[0,1],packag:2,page:2,paper:[0,1],param:1,paramet:[1,2],pass:1,path:[0,1],patienc:[0,1],pdf:0,per:[0,1],percentag:[0,1],perform:[0,1],peter:1,pfister:0,pin:1,pin_memori:1,pip:0,place:0,plot:0,poetri:0,point:1,posit:1,post:1,post_embed_dim:1,pre:2,pred:0,predict:[0,1],predict_func:1,predict_proba:1,predictdataset:1,prefix:1,prepar:1,prepare_target:1,preprint:0,pretrain:[0,2],pretraining_exampl:0,pretraining_ratio:[0,1],pretraining_util:2,previous:1,print:1,prior:1,probabl:1,problem:[1,2],process:1,processed_feat:1,product:1,propos:0,provid:1,pytorch:[1,2],pytorch_tabnet:0,qualifi:1,question:0,rais:1,random:[0,1],randomobfusc:1,rang:0,rapidli:1,rate:0,readm:2,realli:0,recip:1,recommend:0,reconstruct:[0,1],record:1,reduc:[0,1],reducing_matrix:1,regist:1,regress:[0,1],rel:1,repositori:0,repres:1,reproduc:0,res:1,reset:1,result:1,retriev:[0,1],reus:0,reusag:0,risk:0,rmse:[0,1],rmsle:[0,1],roc_auc_scor:0,root:1,run:[0,1],same:[0,1],sampl:[0,1],sampler:1,save:[0,1],save_model:1,saving_path:0,scale:1,schedul:[0,1],scheduler_fn:[0,1],scheduler_param:[0,1],scikit:[0,1],score:[0,1],search:2,section:0,see:[0,1],seed:[0,1],select:0,self:[0,1],semi:2,separ:1,sequenc:1,serializ:1,set:[0,1],set_param:1,set_train:1,shape:1,share:0,shared_lay:1,should:[0,1],show:1,silent:1,simpl:0,sinc:[0,1],singl:[0,1],size:[0,1],skipkei:1,sklearn:[0,1],slack:0,small:1,smaller:1,softmax:1,solut:0,sort:1,sort_kei:1,sourc:[1,2],spars:1,sparsemax:[0,2],sparsemaxfunct:1,sparser:0,sparsiti:0,specif:[0,1],specifii:0,spin:1,squar:1,stabl:1,stack_batch:1,start:[0,1],step:[0,1],step_siz:0,steplr:0,steps_output:1,stop:[0,1],store:1,str:[0,1],string:[0,1],subclass:1,subprocess:1,subsampl:1,sum:1,supermodul:1,supervis:[1,2],support:1,sure:1,tab_model:[0,2],tab_network:2,tabmodel:1,tabnet:[1,2],tabnetclassifi:[0,1],tabnetdecod:1,tabnetencod:1,tabnetmultitaskclassifi:[0,1],tabnetnoembed:1,tabnetpretrain:[0,1],tabnetregressor:[0,1],tabular:2,take:1,talk:0,target:[0,1],target_mapp:1,target_typ:1,task:[0,1],tasks_dim:1,tasks_label:1,tensor:1,term:0,termin:[0,1],than:[0,1],thei:0,them:1,thi:[0,1],tol:1,torch:[0,1],torchdataset:1,train:[1,2],train_dataload:1,train_label:1,trainer:1,trainng:0,transform:[0,1],trick:1,tupl:[0,1],two:1,type:[0,1],type_of_target:1,typeerror:1,typic:0,uniqu:[0,1],unique_label:1,unit:0,unknown:1,unsupervis:1,unsupervised_model:[0,1],unsupervisedloss:1,unsupervisedmetr:1,unsupmetriccontain:1,untouch:0,updat:1,update_fit_param:1,use:[1,2],used:[0,1],user:1,using:0,usual:0,util:[0,2],val_metr:1,valid:[0,1],valid_dataload:1,validate_eval_set:1,valu:[0,1],valueerror:1,variabl:1,vector:1,verbos:[0,1],via:0,video:0,virtual:[],virtual_batch_s:[0,1],vlad:1,wait:1,wan:0,want:0,weight:[0,1],weight_updat:1,were:1,what:2,when:[0,1],where:[0,1],whether:[0,1],which:1,width:0,wihtout:[],within:[0,1],without:[0,1],worker:[0,1],wors:0,wrapper:1,wrong:1,www:[],x_predict:0,x_test:0,x_train:[0,1],x_valid:0,y_pred:1,y_score:[0,1],y_train:[0,1],y_true:[0,1],y_valid:0,you:[0,1],your:0,youtu:[],ysbazo8ymx8:[],zerodivisionerror:1,zip:1},titles:["README","pytorch_tabnet package","Welcome to pytorch_tabnet\u2019s documentation!"],titleterms:{"class":1,"default":0,"function":1,Useful:0,abstract_model:1,attent:0,callback:1,code:0,cpu:0,custom:0,doc:[],document:2,doe:0,early_stopping_metr:[],easi:0,eval_metr:0,evalu:0,fit:0,gpu:0,handl:0,how:0,indic:2,instal:0,interpret:0,label:1,learn:0,link:0,metric:[0,1],model:0,modul:1,multi:1,multiclass_util:1,multitask:1,onli:0,packag:1,paramet:0,pre:0,pretrain:1,pretraining_util:1,problem:0,pytorch:0,pytorch_tabnet:[1,2],readm:0,script:[],semi:0,sourc:0,sparsemax:1,supervis:0,tab_model:1,tab_network:1,tabl:2,tabnet:0,tabular:0,train:0,use:0,util:1,welcom:2,what:0}})
\ No newline at end of file
+Search.setIndex({docnames:["generated_docs/README","generated_docs/pytorch_tabnet","index"],envversion:{"sphinx.domains.c":1,"sphinx.domains.changeset":1,"sphinx.domains.citation":1,"sphinx.domains.cpp":1,"sphinx.domains.index":1,"sphinx.domains.javascript":1,"sphinx.domains.math":2,"sphinx.domains.python":1,"sphinx.domains.rst":1,"sphinx.domains.std":1,"sphinx.ext.viewcode":1,sphinx:56},filenames:["generated_docs/README.md","generated_docs/pytorch_tabnet.rst","index.rst"],objects:{"pytorch_tabnet.abstract_model":{TabModel:[1,1,1,""]},"pytorch_tabnet.abstract_model.TabModel":{cat_dims:[1,2,1,""],cat_emb_dim:[1,2,1,""],cat_idxs:[1,2,1,""],clip_value:[1,2,1,""],compute_loss:[1,3,1,""],device_name:[1,2,1,""],epsilon:[1,2,1,""],explain:[1,3,1,""],fit:[1,3,1,""],gamma:[1,2,1,""],input_dim:[1,2,1,""],lambda_sparse:[1,2,1,""],load_class_attrs:[1,3,1,""],load_model:[1,3,1,""],load_weights_from_unsupervised:[1,3,1,""],mask_type:[1,2,1,""],momentum:[1,2,1,""],n_a:[1,2,1,""],n_d:[1,2,1,""],n_indep_decoder:[1,2,1,""],n_independent:[1,2,1,""],n_shared:[1,2,1,""],n_shared_decoder:[1,2,1,""],n_steps:[1,2,1,""],optimizer_fn:[1,2,1,""],optimizer_params:[1,2,1,""],output_dim:[1,2,1,""],predict:[1,3,1,""],prepare_target:[1,3,1,""],save_model:[1,3,1,""],scheduler_fn:[1,2,1,""],scheduler_params:[1,2,1,""],seed:[1,2,1,""],update_fit_params:[1,3,1,""],verbose:[1,2,1,""]},"pytorch_tabnet.augmentations":{ClassificationSMOTE:[1,1,1,""],RegressionSMOTE:[1,1,1,""]},"pytorch_tabnet.callbacks":{Callback:[1,1,1,""],CallbackContainer:[1,1,1,""],EarlyStopping:[1,1,1,""],History:[1,1,1,""],LRSchedulerCallback:[1,1,1,""]},"pytorch_tabnet.callbacks.Callback":{on_batch_begin:[1,3,1,""],on_batch_end:[1,3,1,""],on_epoch_begin:[1,3,1,""],on_epoch_end:[1,3,1,""],on_train_begin:[1,3,1,""],on_train_end:[1,3,1,""],set_params:[1,3,1,""],set_trainer:[1,3,1,""]},"pytorch_tabnet.callbacks.CallbackContainer":{append:[1,3,1,""],callbacks:[1,2,1,""],on_batch_begin:[1,3,1,""],on_batch_end:[1,3,1,""],on_epoch_begin:[1,3,1,""],on_epoch_end:[1,3,1,""],on_train_begin:[1,3,1,""],on_train_end:[1,3,1,""],set_params:[1,3,1,""],set_trainer:[1,3,1,""]},"pytorch_tabnet.callbacks.EarlyStopping":{early_stopping_metric:[1,2,1,""],is_maximize:[1,2,1,""],on_epoch_end:[1,3,1,""],on_train_end:[1,3,1,""],patience:[1,2,1,""],tol:[1,2,1,""]},"pytorch_tabnet.callbacks.History":{on_batch_end:[1,3,1,""],on_epoch_begin:[1,3,1,""],on_epoch_end:[1,3,1,""],on_train_begin:[1,3,1,""],trainer:[1,2,1,""],verbose:[1,2,1,""]},"pytorch_tabnet.callbacks.LRSchedulerCallback":{early_stopping_metric:[1,2,1,""],is_batch_level:[1,2,1,""],on_batch_end:[1,3,1,""],on_epoch_end:[1,3,1,""],optimizer:[1,2,1,""],scheduler_fn:[1,2,1,""],scheduler_params:[1,2,1,""]},"pytorch_tabnet.metrics":{AUC:[1,1,1,""],Accuracy:[1,1,1,""],BalancedAccuracy:[1,1,1,""],LogLoss:[1,1,1,""],MAE:[1,1,1,""],MSE:[1,1,1,""],Metric:[1,1,1,""],MetricContainer:[1,1,1,""],RMSE:[1,1,1,""],RMSLE:[1,1,1,""],UnsupMetricContainer:[1,1,1,""],UnsupervisedLoss:[1,4,1,""],UnsupervisedLossNumpy:[1,4,1,""],UnsupervisedMetric:[1,1,1,""],UnsupervisedNumpyMetric:[1,1,1,""],check_metrics:[1,4,1,""]},"pytorch_tabnet.metrics.Metric":{get_metrics_by_names:[1,3,1,""]},"pytorch_tabnet.metrics.MetricContainer":{metric_names:[1,2,1,""],prefix:[1,2,1,""]},"pytorch_tabnet.metrics.UnsupMetricContainer":{metric_names:[1,2,1,""],prefix:[1,2,1,""]},"pytorch_tabnet.multiclass_utils":{assert_all_finite:[1,4,1,""],check_classification_targets:[1,4,1,""],check_output_dim:[1,4,1,""],check_unique_type:[1,4,1,""],infer_multitask_output:[1,4,1,""],infer_output_dim:[1,4,1,""],is_multilabel:[1,4,1,""],type_of_target:[1,4,1,""],unique_labels:[1,4,1,""]},"pytorch_tabnet.multitask":{TabNetMultiTaskClassifier:[1,1,1,""]},"pytorch_tabnet.multitask.TabNetMultiTaskClassifier":{cat_dims:[1,2,1,""],cat_idxs:[1,2,1,""],compute_loss:[1,3,1,""],optimizer_params:[1,2,1,""],predict:[1,3,1,""],predict_proba:[1,3,1,""],prepare_target:[1,3,1,""],scheduler_params:[1,2,1,""],stack_batches:[1,3,1,""],update_fit_params:[1,3,1,""]},"pytorch_tabnet.pretraining":{TabNetPretrainer:[1,1,1,""]},"pytorch_tabnet.pretraining.TabNetPretrainer":{cat_dims:[1,2,1,""],cat_idxs:[1,2,1,""],compute_loss:[1,3,1,""],fit:[1,3,1,""],optimizer_params:[1,2,1,""],predict:[1,3,1,""],prepare_target:[1,3,1,""],scheduler_params:[1,2,1,""],stack_batches:[1,3,1,""],update_fit_params:[1,3,1,""]},"pytorch_tabnet.pretraining_utils":{create_dataloaders:[1,4,1,""],validate_eval_set:[1,4,1,""]},"pytorch_tabnet.sparsemax":{Entmax15:[1,1,1,""],Entmax15Function:[1,1,1,""],Entmoid15:[1,1,1,""],Sparsemax:[1,1,1,""],SparsemaxFunction:[1,1,1,""],entmax15:[1,4,1,""],entmoid15:[1,4,1,""],sparsemax:[1,4,1,""]},"pytorch_tabnet.sparsemax.Entmax15":{forward:[1,3,1,""]},"pytorch_tabnet.sparsemax.Entmax15Function":{backward:[1,3,1,""],forward:[1,3,1,""]},"pytorch_tabnet.sparsemax.Entmoid15":{backward:[1,3,1,""],forward:[1,3,1,""]},"pytorch_tabnet.sparsemax.Sparsemax":{forward:[1,3,1,""]},"pytorch_tabnet.sparsemax.SparsemaxFunction":{backward:[1,3,1,""],forward:[1,3,1,""]},"pytorch_tabnet.tab_model":{TabNetClassifier:[1,1,1,""],TabNetRegressor:[1,1,1,""]},"pytorch_tabnet.tab_model.TabNetClassifier":{cat_dims:[1,2,1,""],cat_idxs:[1,2,1,""],compute_loss:[1,3,1,""],optimizer_params:[1,2,1,""],predict_func:[1,3,1,""],predict_proba:[1,3,1,""],prepare_target:[1,3,1,""],scheduler_params:[1,2,1,""],stack_batches:[1,3,1,""],update_fit_params:[1,3,1,""],weight_updater:[1,3,1,""]},"pytorch_tabnet.tab_model.TabNetRegressor":{cat_dims:[1,2,1,""],cat_idxs:[1,2,1,""],compute_loss:[1,3,1,""],optimizer_params:[1,2,1,""],predict_func:[1,3,1,""],prepare_target:[1,3,1,""],scheduler_params:[1,2,1,""],stack_batches:[1,3,1,""],update_fit_params:[1,3,1,""]},"pytorch_tabnet.tab_network":{AttentiveTransformer:[1,1,1,""],EmbeddingGenerator:[1,1,1,""],FeatTransformer:[1,1,1,""],GBN:[1,1,1,""],GLU_Block:[1,1,1,""],GLU_Layer:[1,1,1,""],RandomObfuscator:[1,1,1,""],TabNet:[1,1,1,""],TabNetDecoder:[1,1,1,""],TabNetEncoder:[1,1,1,""],TabNetNoEmbeddings:[1,1,1,""],TabNetPretraining:[1,1,1,""],initialize_glu:[1,4,1,""],initialize_non_glu:[1,4,1,""]},"pytorch_tabnet.tab_network.AttentiveTransformer":{forward:[1,3,1,""]},"pytorch_tabnet.tab_network.EmbeddingGenerator":{forward:[1,3,1,""]},"pytorch_tabnet.tab_network.FeatTransformer":{forward:[1,3,1,""]},"pytorch_tabnet.tab_network.GBN":{forward:[1,3,1,""]},"pytorch_tabnet.tab_network.GLU_Block":{forward:[1,3,1,""]},"pytorch_tabnet.tab_network.GLU_Layer":{forward:[1,3,1,""]},"pytorch_tabnet.tab_network.RandomObfuscator":{forward:[1,3,1,""]},"pytorch_tabnet.tab_network.TabNet":{forward:[1,3,1,""],forward_masks:[1,3,1,""]},"pytorch_tabnet.tab_network.TabNetDecoder":{forward:[1,3,1,""]},"pytorch_tabnet.tab_network.TabNetEncoder":{forward:[1,3,1,""],forward_masks:[1,3,1,""]},"pytorch_tabnet.tab_network.TabNetNoEmbeddings":{forward:[1,3,1,""],forward_masks:[1,3,1,""]},"pytorch_tabnet.tab_network.TabNetPretraining":{forward:[1,3,1,""],forward_masks:[1,3,1,""]},"pytorch_tabnet.utils":{ComplexEncoder:[1,1,1,""],PredictDataset:[1,1,1,""],TorchDataset:[1,1,1,""],check_input:[1,4,1,""],check_warm_start:[1,4,1,""],create_dataloaders:[1,4,1,""],create_explain_matrix:[1,4,1,""],create_sampler:[1,4,1,""],define_device:[1,4,1,""],filter_weights:[1,4,1,""],validate_eval_set:[1,4,1,""]},"pytorch_tabnet.utils.ComplexEncoder":{"default":[1,3,1,""]},pytorch_tabnet:{abstract_model:[1,0,0,"-"],augmentations:[1,0,0,"-"],callbacks:[1,0,0,"-"],metrics:[1,0,0,"-"],multiclass_utils:[1,0,0,"-"],multitask:[1,0,0,"-"],pretraining:[1,0,0,"-"],pretraining_utils:[1,0,0,"-"],sparsemax:[1,0,0,"-"],tab_model:[1,0,0,"-"],tab_network:[1,0,0,"-"],utils:[1,0,0,"-"]}},objnames:{"0":["py","module","Python module"],"1":["py","class","Python class"],"2":["py","attribute","Python attribute"],"3":["py","method","Python method"],"4":["py","function","Python function"]},objtypes:{"0":"py:module","1":"py:class","2":"py:attribute","3":"py:method","4":"py:function"},terms:{"1st":0,"abstract":1,"boolean":1,"case":1,"class":0,"default":[1,2],"float":[0,1],"function":0,"import":[0,1],"int":[0,1],"new":[0,1],"return":[0,1],"static":1,"throw":1,"true":[0,1],"try":[0,1],"while":1,Added:0,For:1,One:1,The:[0,1],Use:1,Useful:2,Using:1,__call__:0,__init__:0,_contextmethodmixin:1,_maxim:0,_name:0,a_max:0,a_min:0,about:1,abov:1,abs:1,absolut:1,abstract_model:2,accept:1,accord:[0,1],accuraci:[0,1],adam:[0,1],added:0,after:1,afterward:1,alia:1,all:[0,1],allow:[0,1],allow_nan:1,along:1,alpha:1,alreadi:0,also:1,although:1,ambigu:1,amount:1,ani:[0,1],anyth:1,api:0,append:1,appli:[0,1],arbitrari:1,architectur:0,argument:1,arik:0,arrai:[0,1],arxiv:[0,1],assert_all_finit:1,assign:0,astudillo:1,attent:[1,2],attentivetransform:1,attribut:1,auc:[0,1],augment:2,auto:[0,1],autograd:1,autom:[0,1],automat:[0,1],avail:0,averag:1,avoid:1,backward:1,balanc:[0,1],balancedaccuraci:1,base:1,baseestim:1,basic:0,batch:[0,1],batch_out:[],batch_siz:[0,1],becaus:[0,1],been:0,befor:[0,1],bellow:0,below:0,ben:1,best:0,beta:1,better:0,between:[0,1],bigger:0,binari:[0,1],blob:0,block:[0,1],bool:[0,1],both:1,build:[0,1],built:1,call:1,callabl:1,callback:[0,2],callbackcontain:1,can:[0,1],capac:0,care:1,cat:[],cat_dim:[0,1],cat_emb_dim:[0,1],cat_idx:[0,1],categor:[0,1],censu:[],certain:1,chang:[0,1],check:1,check_circular:1,check_classification_target:1,check_input:1,check_metr:1,check_nan:[],check_output_dim:1,check_unique_typ:1,check_warm_start:1,choic:0,cite:1,class_attr:1,classic:[0,1],classif:[0,1],classificationsmot:[0,1],classmethod:1,clear:1,clf:0,clip:[0,1],clip_valu:[0,1],clone:0,close:0,cls:1,code:2,coeffici:0,column:1,com:[0,1],commit:0,compat:[0,1],complet:0,complexencod:1,comput:1,compute_loss:1,conda:0,consecut:[0,1],contain:[0,1],content:2,context:1,continu:1,contribut:2,convert:1,corr:1,correct:1,correl:0,correspond:1,could:[0,1],counter:1,cpu:1,creat:[0,1],create_dataload:1,create_explain_matrix:1,create_sampl:1,cross:0,ctx:1,cuda:1,current:[0,1],custom:[1,2],data:[1,2],datafram:1,dataload:[0,1],dataset:1,dblp:1,decai:0,decis:0,decod:0,deduc:0,deep:1,deeprecomodel:1,def:[0,1],defin:[0,1],define_devic:1,degener:1,depend:[0,1],descript:1,detail:1,detect:[0,1],determin:1,develop:0,devic:1,device_nam:[0,1],dict:[0,1],dictionari:1,dictionnari:[0,1],did:1,differ:1,differenti:1,difficulti:0,dim:1,dimens:1,discret:1,discuss:0,disk:0,distinct:1,distribut:1,divid:0,divis:1,docker:0,doe:[1,2],doing:0,don:1,dreamquark:0,dreamquarktabnet:0,drop:[0,1],drop_last:[0,1],dure:[0,1],each:[0,1],earli:[0,1],early_stopping_metr:1,earlystop:1,easi:2,easier:0,easili:0,either:[0,1],element:1,els:1,emb:[],embded:[],embed:[0,1],embedded_x:1,embeddinggener:1,enabl:0,encod:1,end:0,ensur:1,ensure_ascii:1,entmax15:1,entmax15funct:1,entmax:[0,1],entmoid15:1,entropi:0,epoch:[0,1],eps:1,epsilon:[0,1],equal:1,equival:1,error:1,eval:[0,1,2],eval_metr:1,eval_nam:[0,1],eval_set:[0,1],evalu:[1,2],event:1,everi:[0,1],exact:1,exampl:[0,1],except:1,exist:0,exit:1,expert:1,explain:1,explan:1,explanatori:0,explicit:1,extra:0,extract:1,factori:1,fals:[0,1],feattransform:1,featu:[],featur:[0,1],few:0,file:1,filepath:1,filter_weight:1,first:[0,1],fit:[1,2],fly:2,follow:[0,1],forest:[],forg:0,format:1,former:1,formula:1,forward:1,forward_mask:1,found:0,frequenc:1,from:[0,1],from_unsupervis:[0,1],further:0,gamma:[0,1],gate:0,gbn:1,gener:1,get:[0,1],get_metrics_by_nam:1,ghost:[0,1],gini:0,git:0,github:[0,1],give:[0,1],given:[0,1],glu:[0,1],glu_block:1,glu_lay:1,good:0,gpu:1,grad_output:1,gradient:[0,1],greater:0,handl:2,happen:1,harder:0,has:[0,1],have:[0,1],help:0,here:0,highli:1,histori:1,hold:1,hook:1,hot:1,how:[1,2],html:1,http:[0,1],idx:[],ignor:1,imeplement:[],implement:[0,1],improv:[0,1],includ:0,incomplet:1,indent:1,independ:[0,1],index:[1,2],indic:[0,1],infer:1,infer_multitask_output:1,infer_output_dim:1,infin:1,initi:[0,1],initialize_glu:1,initialize_non_glu:1,input:[0,1],input_dim:1,insid:0,instal:2,instanc:1,instead:1,integ:[0,1],interpret:2,invers:[0,1],ipynb:0,is_batch_level:1,is_maxim:1,is_multilabel:1,issu:0,iter:1,its:0,join:0,journal:1,json:1,jsonencod:1,jupyt:0,kaggl:0,keep:1,kei:[0,1],labda:[],lambda:1,lambda_spars:[0,1],larg:0,last:[0,1],later:0,latter:1,layer:0,learn:[1,2],least:[0,1],left:[0,1],length:[0,1],let:1,like:1,line:0,linear:0,link:2,list:[0,1],list_embedded_x:1,list_obfusc:1,list_output:1,list_y_scor:1,list_y_tru:1,load:[1,2],load_class_attr:1,load_model:[0,1],load_weights_from_unsupervis:1,loaded_clf:0,local:[0,1],log:1,logarithm:1,logloss:[0,1],longtensor:1,loop:1,loss:[0,1],loss_fn:[0,1],lower:0,lr_schedul:[0,1],lrschedulercallback:1,m_explain:1,mae:[0,1],main:1,make:[0,1],mandatori:0,mani:1,manual:1,map:1,martin:1,martinsa16:1,mask:[0,1],mask_typ:[0,1],match:0,matric:1,matrix:1,max:0,max_epoch:[0,1],maxim:[0,1],maximum:[0,1],mean:[0,1],mean_squared_log_error:1,memori:1,mention:0,method:1,metric:2,metric_nam:1,metriccontain:1,might:[0,1],mini:0,minimum:1,mix:1,moa:0,modal:0,model:[1,2],model_nam:0,modul:2,moment:1,momentum:[0,1],monitor:1,more:[0,1],most:1,mse:[0,1],multi:0,multiclass:[0,1],multiclass_util:2,multilabel:1,multioutput:1,multipl:1,multitask:[0,2],must:[0,1],n_a:[0,1],n_d:[0,1],n_glu:1,n_glu_independ:1,n_indep_decod:[0,1],n_independ:[0,1],n_sampl:1,n_share:[0,1],n_shared_decod:[0,1],n_step:[0,1],n_unique_label:1,name:[0,1],nan:1,ndarrai:1,need:[0,1],needs_input_grad:1,neg:[0,1],network:1,neural:1,nicula:1,non:1,none:[0,1],normal:[0,1],note:[0,1],notebook:0,now:0,num:[],num_work:[0,1],number:[0,1],numpi:1,obf_var:1,obfusc:1,obj:1,object:1,occur:[],occurr:0,on_batch_begin:1,on_batch_end:1,on_epoch_begin:1,on_epoch_end:1,on_train_begin:1,on_train_end:1,one:[0,1],onecyclelr:1,ones:0,onli:1,oper:1,optim:[0,1],optimizer_fn:[0,1],optimizer_param:[0,1],optimo:[],option:0,order:[0,1],org:[0,1],orgin:1,origin:[0,1],other:1,otherwis:1,our:0,out:1,output:1,output_dim:1,over:1,overfit:0,overridden:1,overwritten:0,own:[0,1],packag:2,page:2,panda:1,paper:[0,1],param:1,paramet:[1,2],pass:1,path:1,patienc:[0,1],pdf:0,per:[0,1],percentag:[0,1],perform:[0,1],peter:1,pfister:0,pin:1,pin_memori:1,pip:0,pipelin:0,place:0,pleas:0,plot:0,poetri:0,point:1,posit:1,possibl:0,post:1,post_embed_dim:1,pre:2,pred:0,predict:[0,1],predict_func:1,predict_proba:1,predictdataset:1,prefix:1,prepar:1,prepare_target:1,preprint:0,present:0,pretrain:[0,2],pretraining_exampl:0,pretraining_ratio:[0,1],pretraining_util:2,previous:1,print:1,prior:1,probabl:1,problem:[1,2],process:1,processed_feat:1,product:[0,1],propos:0,provid:1,pytorch:[1,2],pytorch_tabnet:0,qualifi:1,question:0,rais:1,random:[0,1],randomobfusc:1,rang:0,rapidli:1,rate:0,readi:0,readm:2,realli:0,recip:1,recommend:0,reconstruct:[0,1],record:1,reduc:[0,1],reducing_matrix:1,regist:1,regress:[0,1],regressionsmot:[0,1],rel:1,repositori:0,repres:1,reproduc:0,res:1,reset:1,result:1,retriev:[0,1],reus:0,reusag:0,risk:0,rmse:[0,1],rmsle:[0,1],roc_auc_scor:0,root:1,row:1,rule:[0,1],run:1,same:[0,1],sampl:[0,1],sampler:1,save:[1,2],save_model:[0,1],saved_filepath:0,saving_path:[],saving_path_nam:0,scale:1,schedul:[0,1],scheduler_fn:[0,1],scheduler_param:[0,1],scikit:[0,1],score:[0,1],search:2,section:[],see:[0,1],seed:[0,1],select:0,self:[0,1],semi:2,separ:1,sequenc:1,serializ:1,set:[0,1],set_param:1,set_train:1,shape:1,share:0,shared_lay:1,should:[0,1],show:1,silent:1,simpl:0,sinc:[0,1],singl:[0,1],size:[0,1],skipkei:1,sklearn:[0,1],slack:0,small:1,smaller:1,smote:1,softmax:1,solut:0,sort:1,sort_kei:1,sourc:[1,2],spars:1,sparsemax:[0,2],sparsemaxfunct:1,sparser:0,sparsiti:0,specif:[0,1],specifii:0,spin:1,squar:1,stabl:1,stack_batch:1,stai:1,start:[0,1],state:0,step:[0,1],step_siz:0,steplr:0,steps_output:1,stop:[0,1],store:1,str:[0,1],string:[0,1],subclass:1,subprocess:1,subsampl:1,sum:1,supermodul:1,supervis:[1,2],support:1,sure:[0,1],tab_model:[0,2],tab_network:2,tabmodel:1,tabnet:[1,2],tabnet_model_test_1:0,tabnetclassifi:[0,1],tabnetdecod:1,tabnetencod:1,tabnetmultitaskclassifi:[0,1],tabnetnoembed:1,tabnetpretrain:[0,1],tabnetregressor:[0,1],tabular:2,take:1,talk:0,target:[0,1],target_mapp:1,target_typ:1,task:[0,1],tasks_dim:1,tasks_label:1,templat:0,tensor:1,term:0,termin:[0,1],than:[0,1],thei:0,them:1,thi:[0,1],tol:1,torch:[0,1],torchdataset:1,train:[1,2],train_dataload:1,train_label:1,trainer:1,trainng:0,transform:[0,1],trick:1,tupl:[0,1],twice:0,two:1,type:[0,1],type_of_target:1,typeerror:1,typic:0,unchang:1,uniqu:[0,1],unique_label:1,unit:0,unknown:1,unsupervis:1,unsupervised_model:[0,1],unsupervisedloss:1,unsupervisedlossnumpi:1,unsupervisedmetr:1,unsupervisednumpymetr:1,unsupmetriccontain:1,untouch:0,updat:1,update_fit_param:1,usag:1,use:[1,2],used:[0,1],useful:0,user:1,using:0,usual:0,util:[0,2],val_metr:1,valid:[0,1],valid_dataload:1,validate_eval_set:1,valu:[0,1],valueerror:1,variabl:1,vector:1,verbos:[0,1],via:0,video:0,virtual:[],virtual_batch_s:[0,1],vlad:1,wait:1,wan:0,want:0,warm:0,warm_start:[0,1],warn:1,weight:[0,1],weight_updat:1,well:1,were:1,what:2,when:[0,1],where:1,wheter:1,whether:[0,1],which:1,width:0,wihtout:[],wish:0,within:[0,1],without:[0,1],work:1,worker:[0,1],wors:0,wrapper:1,wrong:1,www:[],x_predict:0,x_test:0,x_train:[0,1],x_valid:0,y_pred:1,y_score:[0,1],y_train:[0,1],y_true:[0,1],y_valid:0,you:[0,1],your:0,youtu:[],ysbazo8ymx8:[],zerodivisionerror:1,zip:1},titles:["README","pytorch_tabnet package","Welcome to pytorch_tabnet\u2019s documentation!"],titleterms:{"class":1,"default":0,"function":1,Useful:0,abstract_model:1,attent:0,augment:[0,1],callback:1,code:0,contribut:0,cpu:0,custom:0,data:0,doc:[],document:2,doe:0,early_stopping_metr:[],easi:0,eval_metr:0,evalu:0,fit:0,fly:0,gpu:0,handl:0,how:0,indic:2,instal:0,interpret:0,label:1,learn:0,link:0,load:0,metric:[0,1],model:0,modul:1,multi:1,multiclass_util:1,multitask:1,onli:0,packag:1,paramet:0,pre:0,pretrain:1,pretraining_util:1,problem:0,pytorch:0,pytorch_tabnet:[1,2],readm:0,save:0,script:[],semi:0,sourc:0,sparsemax:1,supervis:0,tab_model:1,tab_network:1,tabl:2,tabnet:0,tabular:0,train:0,use:0,util:1,welcom:2,what:0}})
\ No newline at end of file
diff --git a/pyproject.toml b/pyproject.toml
index 1f96b264..f9fa0916 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "pytorch_tabnet"
-version = "3.1.1"
+version = "4.0"
 description = "PyTorch implementation of TabNet"
 homepage = "https://github.com/dreamquark-ai/tabnet"
 repository = "https://github.com/dreamquark-ai/tabnet"