Repository: jrzaurin/pytorch-widedeep (language: python)

Issue:
Pmulinka/raytune
@jrzaurin please take a look at it this draft, main points, and logic:
* I added RayTuneReporter callback that reports history and possibly lr_hisotry if used in callbacks
  * I do it by sorting callbacks and running RayTuneReporter as the last one that checks history/lr_hisotry and sends/reports it to raytune
  * you can check it yourself in examples notebook num.12
* I adjusted how the history is saved, ie. if the array of values is saved I separate it into columns,eg. if you have binary classification and you want to see a recall for both classes (my current use case in imbalanced dataset scenario) it separates it into recall_0 and recall_1
* I adjusted tab2vec function as I am testing learned entity embeddings from the deep nets in the lightgbm regressor, for this case it can be useful to have pandas dataframe as an output from tab2vec to have some idea about feature names for follow-up analysis by shapely or other framework + there is a boolean switch to get the list of new embedded columns names - again this might be useful as if you do entity embedding of categorical features in lightgbm using embedding learned from deep-net you don't want to scale these embedded columns (I don't have good enough reason right now , just that it does not sound right :) )


Discussion / hints:

Fix patch:
diff --git a/VERSION b/VERSION
index 437d26b1..86844988 100644
--- a/VERSION
+++ b/VERSION
@@ -1,1 +1,1 @@
-1.0.10
\ No newline at end of file
+1.0.11
\ No newline at end of file
diff --git a/docs/callbacks.rst b/docs/callbacks.rst
index c5ef2ef2..2ba47533 100644
--- a/docs/callbacks.rst
+++ b/docs/callbacks.rst
@@ -24,3 +24,6 @@ Here are the 4 callbacks available in ``pytorch-widedepp``: ``History``,
 
 .. autoclass:: pytorch_widedeep.callbacks.EarlyStopping
 	:members:
+
+.. autoclass:: pytorch_widedeep.callbacks.RayTuneReporter
+	:members:
diff --git a/examples/01_Preprocessors_and_utils.ipynb b/examples/01_Preprocessors_and_utils.ipynb
index 25439b46..7abbca0c 100644
--- a/examples/01_Preprocessors_and_utils.ipynb
+++ b/examples/01_Preprocessors_and_utils.ipynb
@@ -262,8 +262,15 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "wide_cols = ['education', 'relationship','workclass','occupation','native-country','gender']\n",
-    "crossed_cols = [('education', 'occupation'), ('native-country', 'occupation')]"
+    "wide_cols = [\n",
+    "    \"education\",\n",
+    "    \"relationship\",\n",
+    "    \"workclass\",\n",
+    "    \"occupation\",\n",
+    "    \"native-country\",\n",
+    "    \"gender\",\n",
+    "]\n",
+    "crossed_cols = [(\"education\", \"occupation\"), (\"native-country\", \"occupation\")]"
    ]
   },
   {
@@ -440,9 +447,15 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "#cat_embed_cols = [(column_name, embed_dim), ...]\n",
-    "cat_embed_cols = [('education',10), ('relationship',8), ('workclass',10), ('occupation',10),('native-country',10)]\n",
-    "continuous_cols = [\"age\",\"hours-per-week\"]"
+    "# cat_embed_cols = [(column_name, embed_dim), ...]\n",
+    "cat_embed_cols = [\n",
+    "    (\"education\", 10),\n",
+    "    (\"relationship\", 8),\n",
+    "    (\"workclass\", 10),\n",
+    "    (\"occupation\", 10),\n",
+    "    (\"native-country\", 10),\n",
+    "]\n",
+    "continuous_cols = [\"age\", \"hours-per-week\"]"
    ]
   },
   {
@@ -451,7 +464,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "tab_preprocessor = TabPreprocessor(embed_cols=cat_embed_cols, continuous_cols=continuous_cols)\n",
+    "tab_preprocessor = TabPreprocessor(\n",
+    "    embed_cols=cat_embed_cols, continuous_cols=continuous_cols\n",
+    ")\n",
     "X_tab = tab_preprocessor.fit_transform(df)\n",
     "# From here on, any new observation can be prepared by simply running `.transform`\n",
     "# new_X_deep = deep_preprocessor.transform(new_df)"
@@ -986,7 +1001,7 @@
     }
    ],
    "source": [
-    "text_preprocessor = TextPreprocessor(text_col='description')\n",
+    "text_preprocessor = TextPreprocessor(text_col=\"description\")\n",
     "X_text = text_preprocessor.fit_transform(df)\n",
     "# From here on, any new observation can be prepared by simply running `.transform`\n",
     "# new_X_text = text_preprocessor.transform(new_df)"
@@ -2163,7 +2178,9 @@
     }
    ],
    "source": [
-    "image_preprocessor = wd.preprocessing.ImagePreprocessor(img_col='id', img_path=\"data/airbnb/property_picture/\")\n",
+    "image_preprocessor = wd.preprocessing.ImagePreprocessor(\n",
+    "    img_col=\"id\", img_path=\"data/airbnb/property_picture/\"\n",
+    ")\n",
     "X_images = image_preprocessor.fit_transform(df)\n",
     "# From here on, any new observation can be prepared by simply running `.transform`\n",
     "# new_X_images = image_preprocessor.transform(new_df)"
@@ -2241,10 +2258,13 @@
    "outputs": [],
    "source": [
     "import matplotlib.pyplot as plt\n",
+    "\n",
     "%matplotlib inline\n",
     "\n",
-    "def show_img(im, figsize=(5,5), ax=None):\n",
-    "    if not ax: fig,ax = plt.subplots(figsize=figsize)\n",
+    "\n",
+    "def show_img(im, figsize=(5, 5), ax=None):\n",
+    "    if not ax:\n",
+    "        fig, ax = plt.subplots(figsize=figsize)\n",
     "    ax.imshow(im)\n",
     "    ax.set_axis_off()\n",
     "    return ax"
@@ -2297,7 +2317,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "spp = wd.utils.SimplePreprocessor(224,224)\n",
+    "spp = wd.utils.SimplePreprocessor(224, 224)\n",
     "prop_resized_imgs = [spp.preprocess(im) for im in prop_imgs]"
    ]
   },
@@ -2320,8 +2340,8 @@
     }
    ],
    "source": [
-    "fig,axes = plt.subplots(2, 5, figsize=(16,5))\n",
-    "for i,im in enumerate(prop_resized_imgs):\n",
+    "fig, axes = plt.subplots(2, 5, figsize=(16, 5))\n",
+    "for i, im in enumerate(prop_resized_imgs):\n",
     "    show_img(im, ax=axes.flat[i])"
    ]
   },
diff --git a/examples/02_1_Model_Components.ipynb b/examples/02_1_Model_Components.ipynb
index 61873299..204fa2b1 100644
--- a/examples/02_1_Model_Components.ipynb
+++ b/examples/02_1_Model_Components.ipynb
@@ -108,7 +108,7 @@
     }
    ],
    "source": [
-    "df = pd.DataFrame({'color': ['r', 'b', 'g'], 'size': ['s', 'n', 'l']})\n",
+    "df = pd.DataFrame({\"color\": [\"r\", \"b\", \"g\"], \"size\": [\"s\", \"n\", \"l\"]})\n",
     "df.head()"
    ]
   },
@@ -125,7 +125,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "obs_0_oh = (np.array([1., 0., 0., 1., 0., 0.])).astype('float32')"
+    "obs_0_oh = (np.array([1.0, 0.0, 0.0, 1.0, 0.0, 0.0])).astype(\"float32\")"
    ]
   },
   {
@@ -141,7 +141,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "obs_0_le = (np.array([0, 3])).astype('int64')"
+    "obs_0_le = (np.array([0, 3])).astype(\"int64\")"
    ]
   },
   {
@@ -169,7 +169,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "emb = nn.Embedding(6, 1) \n",
+    "emb = nn.Embedding(6, 1)\n",
     "emb.weight = nn.Parameter(lin.weight.reshape_as(emb.weight))"
    ]
   },
@@ -235,7 +235,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "?Wide"
+    "# ?Wide"
    ]
   },
   {
@@ -299,7 +299,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "?DeepText"
+    "# ?DeepText"
    ]
   },
   {
@@ -308,7 +308,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "X_text = torch.cat((torch.zeros([5,1]), torch.empty(5, 4).random_(1,4)), axis=1)"
+    "X_text = torch.cat((torch.zeros([5, 1]), torch.empty(5, 4).random_(1, 4)), axis=1)"
    ]
   },
   {
@@ -365,7 +365,14 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "deeptext = DeepText(vocab_size=4, hidden_dim=8, n_layers=1, padding_idx=0, embed_dim=4, head_hidden_dims=[8,4])"
+    "deeptext = DeepText(\n",
+    "    vocab_size=4,\n",
+    "    hidden_dim=8,\n",
+    "    n_layers=1,\n",
+    "    padding_idx=0,\n",
+    "    embed_dim=4,\n",
+    "    head_hidden_dims=[8, 4],\n",
+    ")"
    ]
   },
   {
@@ -430,7 +437,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "?DeepImage"
+    "# ?DeepImage"
    ]
   },
   {
@@ -439,7 +446,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "X_img = torch.rand((2,3,224,224))"
+    "X_img = torch.rand((2, 3, 224, 224))"
    ]
   },
   {
@@ -695,7 +702,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "?WideDeep"
+    "# ?WideDeep"
    ]
   }
  ],
diff --git a/examples/02_2_deeptabular_models.ipynb b/examples/02_2_deeptabular_models.ipynb
index 37b439ee..495e63a7 100644
--- a/examples/02_2_deeptabular_models.ipynb
+++ b/examples/02_2_deeptabular_models.ipynb
@@ -62,7 +62,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "?TabMlp"
+    "# ?TabMlp"
    ]
   },
   {
@@ -78,12 +78,17 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "colnames = ['a', 'b', 'c', 'd', 'e']\n",
+    "colnames = [\"a\", \"b\", \"c\", \"d\", \"e\"]\n",
     "X_tab = torch.cat((torch.empty(5, 4).random_(4), torch.rand(5, 1)), axis=1)\n",
-    "embed_input = [(u,i,j) for u,i,j in zip(colnames[:4], [4]*4, [8]*4)]\n",
-    "column_idx = {k:v for v,k in enumerate(colnames)}\n",
-    "tabmlp = TabMlp(mlp_hidden_dims=[8,4], continuous_cols=['e'], column_idx=column_idx, \n",
-    "                embed_input=embed_input, cont_norm_layer=\"batchnorm\")\n",
+    "embed_input = [(u, i, j) for u, i, j in zip(colnames[:4], [4] * 4, [8] * 4)]\n",
+    "column_idx = {k: v for v, k in enumerate(colnames)}\n",
+    "tabmlp = TabMlp(\n",
+    "    mlp_hidden_dims=[8, 4],\n",
+    "    continuous_cols=[\"e\"],\n",
+    "    column_idx=column_idx,\n",
+    "    embed_input=embed_input,\n",
+    "    cont_norm_layer=\"batchnorm\",\n",
+    ")\n",
     "out = tabmlp(X_tab)"
    ]
   },
@@ -254,7 +259,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "?TabResnet"
+    "# ?TabResnet"
    ]
   },
   {
@@ -264,17 +269,19 @@
    "outputs": [],
    "source": [
     "X_tab = torch.cat((torch.empty(5, 4).random_(4), torch.rand(5, 1)), axis=1)\n",
-    "colnames = ['a', 'b', 'c', 'd', 'e']\n",
-    "embed_input = [(u,i,j) for u,i,j in zip(colnames[:4], [4]*4, [8]*4)]\n",
-    "column_idx = {k:v for v,k in enumerate(colnames)}\n",
-    "tabresnet = TabResnet(blocks_dims=[16,16,16], \n",
-    "                  column_idx=column_idx, \n",
-    "                  embed_input=embed_input,\n",
-    "                  continuous_cols = ['e'],\n",
-    "                  cont_norm_layer = \"layernorm\",\n",
-    "                  concat_cont_first = False, \n",
-    "                  mlp_hidden_dims = [16, 4],\n",
-    "                  mlp_dropout = 0.5)\n",
+    "colnames = [\"a\", \"b\", \"c\", \"d\", \"e\"]\n",
+    "embed_input = [(u, i, j) for u, i, j in zip(colnames[:4], [4] * 4, [8] * 4)]\n",
+    "column_idx = {k: v for v, k in enumerate(colnames)}\n",
+    "tabresnet = TabResnet(\n",
+    "    blocks_dims=[16, 16, 16],\n",
+    "    column_idx=column_idx,\n",
+    "    embed_input=embed_input,\n",
+    "    continuous_cols=[\"e\"],\n",
+    "    cont_norm_layer=\"layernorm\",\n",
+    "    concat_cont_first=False,\n",
+    "    mlp_hidden_dims=[16, 4],\n",
+    "    mlp_dropout=0.5,\n",
+    ")\n",
     "out = tabresnet(X_tab)"
    ]
   },
@@ -391,7 +398,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "?TabNet"
+    "# ?TabNet"
    ]
   },
   {
@@ -400,18 +407,17 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "\n",
     "X_tab = torch.cat((torch.empty(5, 4).random_(4), torch.rand(5, 1)), axis=1)\n",
-    "colnames = ['a', 'b', 'c', 'd', 'e']\n",
-    "embed_input = [(u,i,j) for u,i,j in zip(colnames[:4], [4]*4, [8]*4)]\n",
-    "column_idx = {k:v for v,k in enumerate(colnames)}\n",
+    "colnames = [\"a\", \"b\", \"c\", \"d\", \"e\"]\n",
+    "embed_input = [(u, i, j) for u, i, j in zip(colnames[:4], [4] * 4, [8] * 4)]\n",
+    "column_idx = {k: v for v, k in enumerate(colnames)}\n",
     "tabnet = TabNet(\n",
-    "        column_idx=column_idx,\n",
-    "        embed_input=embed_input,\n",
-    "        continuous_cols=['e'],\n",
-    "        cont_norm_layer = \"batchnorm\",\n",
-    "        ghost_bn = False,\n",
-    "    )\n",
+    "    column_idx=column_idx,\n",
+    "    embed_input=embed_input,\n",
+    "    continuous_cols=[\"e\"],\n",
+    "    cont_norm_layer=\"batchnorm\",\n",
+    "    ghost_bn=False,\n",
+    ")\n",
     "out = tabnet(X_tab)"
    ]
   },
@@ -645,7 +651,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "?TabTransformer"
+    "# ?TabTransformer"
    ]
   },
   {
@@ -655,12 +661,14 @@
    "outputs": [],
    "source": [
     "X_tab = torch.cat((torch.empty(5, 4).random_(4), torch.rand(5, 1)), axis=1)\n",
-    "colnames = ['a', 'b', 'c', 'd', 'e']\n",
-    "embed_input = [(u,i) for u,i in zip(colnames[:4], [4]*4)]\n",
-    "continuous_cols = ['e']\n",
-    "column_idx = {k:v for v,k in enumerate(colnames)}\n",
-    "tab_transformer = TabTransformer(column_idx=column_idx, embed_input=embed_input, continuous_cols=continuous_cols)\n",
-    "out = tab_transformer(X_tab) "
+    "colnames = [\"a\", \"b\", \"c\", \"d\", \"e\"]\n",
+    "embed_input = [(u, i) for u, i in zip(colnames[:4], [4] * 4)]\n",
+    "continuous_cols = [\"e\"]\n",
+    "column_idx = {k: v for v, k in enumerate(colnames)}\n",
+    "tab_transformer = TabTransformer(\n",
+    "    column_idx=column_idx, embed_input=embed_input, continuous_cols=continuous_cols\n",
+    ")\n",
+    "out = tab_transformer(X_tab)"
    ]
   },
   {
@@ -804,13 +812,13 @@
    "outputs": [],
    "source": [
     "tab_transformer = TabTransformer(\n",
-    "    column_idx=column_idx, \n",
-    "    embed_input=embed_input, \n",
+    "    column_idx=column_idx,\n",
+    "    embed_input=embed_input,\n",
     "    continuous_cols=continuous_cols,\n",
     "    embed_continuous=True,\n",
     "    embed_continuous_activation=\"relu\",\n",
     ")\n",
-    "out = tab_transformer(X_tab) "
+    "out = tab_transformer(X_tab)"
    ]
   },
   {
@@ -992,7 +1000,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "?SAINT"
+    "# ?SAINT"
    ]
   },
   {
@@ -1002,17 +1010,17 @@
    "outputs": [],
    "source": [
     "X_tab = torch.cat((torch.empty(5, 4).random_(4), torch.rand(5, 1)), axis=1)\n",
-    "colnames = ['a', 'b', 'c', 'd', 'e']\n",
-    "embed_input = [(u,i) for u,i in zip(colnames[:4], [4]*4)]\n",
-    "continuous_cols = ['e']\n",
-    "column_idx = {k:v for v,k in enumerate(colnames)}\n",
+    "colnames = [\"a\", \"b\", \"c\", \"d\", \"e\"]\n",
+    "embed_input = [(u, i) for u, i in zip(colnames[:4], [4] * 4)]\n",
+    "continuous_cols = [\"e\"]\n",
+    "column_idx = {k: v for v, k in enumerate(colnames)}\n",
     "saint = SAINT(\n",
-    "    column_idx=column_idx, \n",
-    "    embed_input=embed_input, \n",
+    "    column_idx=column_idx,\n",
+    "    embed_input=embed_input,\n",
     "    continuous_cols=continuous_cols,\n",
     "    embed_continuous_activation=\"leaky_relu\",\n",
     ")\n",
-    "out = saint(X_tab) "
+    "out = saint(X_tab)"
    ]
   },
   {
diff --git a/examples/03_Binary_Classification_with_Defaults.ipynb b/examples/03_Binary_Classification_with_Defaults.ipynb
index 6ed9bebf..7ea151d2 100644
--- a/examples/03_Binary_Classification_with_Defaults.ipynb
+++ b/examples/03_Binary_Classification_with_Defaults.ipynb
@@ -201,7 +201,7 @@
     }
    ],
    "source": [
-    "df = pd.read_csv('data/adult/adult.csv.zip')\n",
+    "df = pd.read_csv(\"data/adult/adult.csv.zip\")\n",
     "df.head()"
    ]
   },
@@ -374,9 +374,9 @@
    "source": [
     "# For convenience, we'll replace '-' with '_'\n",
     "df.columns = [c.replace(\"-\", \"_\") for c in df.columns]\n",
-    "#binary target\n",
-    "df['income_label'] = (df[\"income\"].apply(lambda x: \">50K\" in x)).astype(int)\n",
-    "df.drop('income', axis=1, inplace=True)\n",
+    "# binary target\n",
+    "df[\"income_label\"] = (df[\"income\"].apply(lambda x: \">50K\" in x)).astype(int)\n",
+    "df.drop(\"income\", axis=1, inplace=True)\n",
     "df.head()"
    ]
   },
@@ -395,11 +395,24 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "wide_cols = ['education', 'relationship','workclass','occupation','native_country','gender']\n",
-    "crossed_cols = [('education', 'occupation'), ('native_country', 'occupation')]\n",
-    "cat_embed_cols = [('education',16), ('relationship',8), ('workclass',16), ('occupation',16),('native_country',16)]\n",
-    "continuous_cols = [\"age\",\"hours_per_week\"]\n",
-    "target_col = 'income_label'"
+    "wide_cols = [\n",
+    "    \"education\",\n",
+    "    \"relationship\",\n",
+    "    \"workclass\",\n",
+    "    \"occupation\",\n",
+    "    \"native_country\",\n",
+    "    \"gender\",\n",
+    "]\n",
+    "crossed_cols = [(\"education\", \"occupation\"), (\"native_country\", \"occupation\")]\n",
+    "cat_embed_cols = [\n",
+    "    (\"education\", 16),\n",
+    "    (\"relationship\", 8),\n",
+    "    (\"workclass\", 16),\n",
+    "    (\"occupation\", 16),\n",
+    "    (\"native_country\", 16),\n",
+    "]\n",
+    "continuous_cols = [\"age\", \"hours_per_week\"]\n",
+    "target_col = \"income_label\""
    ]
   },
   {
@@ -416,7 +429,9 @@
     "X_wide = wide_preprocessor.fit_transform(df)\n",
     "\n",
     "# deeptabular\n",
-    "tab_preprocessor = TabPreprocessor(embed_cols=cat_embed_cols, continuous_cols=continuous_cols)\n",
+    "tab_preprocessor = TabPreprocessor(\n",
+    "    embed_cols=cat_embed_cols, continuous_cols=continuous_cols\n",
+    ")\n",
     "X_tab = tab_preprocessor.fit_transform(df)"
    ]
   },
@@ -490,10 +505,12 @@
    "outputs": [],
    "source": [
     "wide = Wide(wide_dim=np.unique(X_wide).shape[0], pred_dim=1)\n",
-    "deeptabular = TabMlp(mlp_hidden_dims=[64,32], \n",
-    "                      column_idx=tab_preprocessor.column_idx,\n",
-    "                      embed_input=tab_preprocessor.embeddings_input,\n",
-    "                      continuous_cols=continuous_cols)\n",
+    "deeptabular = TabMlp(\n",
+    "    mlp_hidden_dims=[64, 32],\n",
+    "    column_idx=tab_preprocessor.column_idx,\n",
+    "    embed_input=tab_preprocessor.embeddings_input,\n",
+    "    continuous_cols=continuous_cols,\n",
+    ")\n",
     "model = WideDeep(wide=wide, deeptabular=deeptabular)"
    ]
   },
@@ -579,7 +596,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "trainer = Trainer(model, objective='binary', metrics=[Accuracy, Precision])"
+    "trainer = Trainer(model, objective=\"binary\", metrics=[Accuracy, Precision])"
    ]
   },
   {
@@ -605,7 +622,9 @@
     }
    ],
    "source": [
-    "trainer.fit(X_wide=X_wide, X_tab=X_tab, target=target, n_epochs=5, batch_size=64, val_split=0.2)"
+    "trainer.fit(\n",
+    "    X_wide=X_wide, X_tab=X_tab, target=target, n_epochs=5, batch_size=64, val_split=0.2\n",
+    ")"
    ]
   },
   {
@@ -624,10 +643,12 @@
    "outputs": [],
    "source": [
     "wide = Wide(wide_dim=np.unique(X_wide).shape[0], pred_dim=1)\n",
-    "deeptabular = TabResnet(blocks_dims=[64,32], \n",
-    "                        column_idx=tab_preprocessor.column_idx,\n",
-    "                        embed_input=tab_preprocessor.embeddings_input,\n",
-    "                        continuous_cols=continuous_cols)\n",
+    "deeptabular = TabResnet(\n",
+    "    blocks_dims=[64, 32],\n",
+    "    column_idx=tab_preprocessor.column_idx,\n",
+    "    embed_input=tab_preprocessor.embeddings_input,\n",
+    "    continuous_cols=continuous_cols,\n",
+    ")\n",
     "model = WideDeep(wide=wide, deeptabular=deeptabular)"
    ]
   },
@@ -637,7 +658,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "trainer = Trainer(model, objective='binary', metrics=[Accuracy, Precision])"
+    "trainer = Trainer(model, objective=\"binary\", metrics=[Accuracy, Precision])"
    ]
   },
   {
@@ -663,7 +684,9 @@
     }
    ],
    "source": [
-    "trainer.fit(X_wide=X_wide, X_tab=X_tab, target=target, n_epochs=5, batch_size=64, val_split=0.2)"
+    "trainer.fit(\n",
+    "    X_wide=X_wide, X_tab=X_tab, target=target, n_epochs=5, batch_size=64, val_split=0.2\n",
+    ")"
    ]
   },
   {
@@ -710,10 +733,12 @@
    "outputs": [],
    "source": [
     "# deeptabular\n",
-    "tab_preprocessor = TabPreprocessor(embed_cols=cat_embed_cols_for_transformer, \n",
-    "                                   continuous_cols=continuous_cols, \n",
-    "                                   for_transformer=True, \n",
-    "                                   with_cls_token = True) # you need to define this since it changes pre-processing\n",
+    "tab_preprocessor = TabPreprocessor(\n",
+    "    embed_cols=cat_embed_cols_for_transformer,\n",
+    "    continuous_cols=continuous_cols,\n",
+    "    for_transformer=True,\n",
+    "    with_cls_token=True,\n",
+    ")  # you need to define this since it changes pre-processing\n",
     "X_tab = tab_preprocessor.fit_transform(df)"
    ]
   },
@@ -724,11 +749,13 @@
    "outputs": [],
    "source": [
     "wide = Wide(wide_dim=np.unique(X_wide).shape[0], pred_dim=1)\n",
-    "deeptabular = TabTransformer(column_idx=tab_preprocessor.column_idx,\n",
-    "                             embed_input=tab_preprocessor.embeddings_input,\n",
-    "                             continuous_cols=continuous_cols, \n",
-    "                             embed_continuous = True, \n",
-    "                             embed_continuous_activation = \"relu\")\n",
+    "deeptabular = TabTransformer(\n",
+    "    column_idx=tab_preprocessor.column_idx,\n",
+    "    embed_input=tab_preprocessor.embeddings_input,\n",
+    "    continuous_cols=continuous_cols,\n",
+    "    embed_continuous=True,\n",
+    "    embed_continuous_activation=\"relu\",\n",
+    ")\n",
     "model = WideDeep(wide=wide, deeptabular=deeptabular)"
    ]
   },
@@ -738,7 +765,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "trainer = Trainer(model, objective='binary', metrics=[Accuracy, Precision])"
+    "trainer = Trainer(model, objective=\"binary\", metrics=[Accuracy, Precision])"
    ]
   },
   {
@@ -758,7 +785,9 @@
     }
    ],
    "source": [
-    "trainer.fit(X_wide=X_wide, X_tab=X_tab, target=target, n_epochs=2, batch_size=512, val_split=0.2)"
+    "trainer.fit(\n",
+    "    X_wide=X_wide, X_tab=X_tab, target=target, n_epochs=2, batch_size=512, val_split=0.2\n",
+    ")"
    ]
   },
   {
@@ -783,7 +812,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "trainer = Trainer(model, objective='binary', metrics=[Accuracy, Precision])"
+    "trainer = Trainer(model, objective=\"binary\", metrics=[Accuracy, Precision])"
    ]
   },
   {
diff --git a/examples/04_Binary_Classification_Varying_Parameters.ipynb b/examples/04_Binary_Classification_Varying_Parameters.ipynb
index 23ba480c..147129d3 100644
--- a/examples/04_Binary_Classification_Varying_Parameters.ipynb
+++ b/examples/04_Binary_Classification_Varying_Parameters.ipynb
@@ -201,7 +201,7 @@
     }
    ],
    "source": [
-    "df = pd.read_csv('data/adult/adult.csv.zip')\n",
+    "df = pd.read_csv(\"data/adult/adult.csv.zip\")\n",
     "df.head()"
    ]
   },
@@ -374,9 +374,9 @@
    "source": [
     "# For convenience, we'll replace '-' with '_'\n",
     "df.columns = [c.replace(\"-\", \"_\") for c in df.columns]\n",
-    "#binary target\n",
-    "df['income_label'] = (df[\"income\"].apply(lambda x: \">50K\" in x)).astype(int)\n",
-    "df.drop('income', axis=1, inplace=True)\n",
+    "# binary target\n",
+    "df[\"income_label\"] = (df[\"income\"].apply(lambda x: \">50K\" in x)).astype(int)\n",
+    "df.drop(\"income\", axis=1, inplace=True)\n",
     "df.head()"
    ]
   },
@@ -395,11 +395,24 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "wide_cols = ['education', 'relationship','workclass','occupation','native_country','gender']\n",
-    "crossed_cols = [('education', 'occupation'), ('native_country', 'occupation')]\n",
-    "cat_embed_cols = [('education',16), ('relationship',8), ('workclass',16), ('occupation',16),('native_country',16)]\n",
-    "continuous_cols = [\"age\",\"hours_per_week\"]\n",
-    "target_col = 'income_label'"
+    "wide_cols = [\n",
+    "    \"education\",\n",
+    "    \"relationship\",\n",
+    "    \"workclass\",\n",
+    "    \"occupation\",\n",
+    "    \"native_country\",\n",
+    "    \"gender\",\n",
+    "]\n",
+    "crossed_cols = [(\"education\", \"occupation\"), (\"native_country\", \"occupation\")]\n",
+    "cat_embed_cols = [\n",
+    "    (\"education\", 16),\n",
+    "    (\"relationship\", 8),\n",
+    "    (\"workclass\", 16),\n",
+    "    (\"occupation\", 16),\n",
+    "    (\"native_country\", 16),\n",
+    "]\n",
+    "continuous_cols = [\"age\", \"hours_per_week\"]\n",
+    "target_col = \"income_label\""
    ]
   },
   {
@@ -416,7 +429,9 @@
     "X_wide = wide_preprocessor.fit_transform(df)\n",
     "\n",
     "# DEEP\n",
-    "tab_preprocessor = TabPreprocessor(embed_cols=cat_embed_cols, continuous_cols=continuous_cols)\n",
+    "tab_preprocessor = TabPreprocessor(\n",
+    "    embed_cols=cat_embed_cols, continuous_cols=continuous_cols\n",
+    ")\n",
     "X_tab = tab_preprocessor.fit_transform(df)"
    ]
   },
@@ -498,7 +513,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "?TabMlp"
+    "# ?TabMlp"
    ]
   },
   {
@@ -509,13 +524,15 @@
    "source": [
     "wide = Wide(wide_dim=np.unique(X_wide).shape[0], pred_dim=1)\n",
     "# We can add dropout and batchnorm to the dense layers, as well as chose the order of the operations\n",
-    "deeptabular = TabMlp(column_idx=tab_preprocessor.column_idx,\n",
-    "                   mlp_hidden_dims=[64,32], \n",
-    "                   mlp_dropout=[0.5, 0.5], \n",
-    "                   mlp_batchnorm=True, \n",
-    "                   mlp_linear_first = True,\n",
-    "                   embed_input=tab_preprocessor.embeddings_input,\n",
-    "                   continuous_cols=continuous_cols)\n",
+    "deeptabular = TabMlp(\n",
+    "    column_idx=tab_preprocessor.column_idx,\n",
+    "    mlp_hidden_dims=[64, 32],\n",
+    "    mlp_dropout=[0.5, 0.5],\n",
+    "    mlp_batchnorm=True,\n",
+    "    mlp_linear_first=True,\n",
+    "    embed_input=tab_preprocessor.embeddings_input,\n",
+    "    continuous_cols=continuous_cols,\n",
+    ")\n",
     "model = WideDeep(wide=wide, deeptabular=deeptabular)"
    ]
   },
@@ -627,11 +644,15 @@
    "outputs": [],
    "source": [
     "# Component-dependent settings as Dict\n",
-    "optimizers = {'wide': wide_opt, 'deeptabular':deep_opt}\n",
-    "schedulers = {'wide': wide_sch, 'deeptabular':deep_sch}\n",
-    "initializers = {'wide': KaimingNormal, 'deeptabular':XavierNormal}\n",
+    "optimizers = {\"wide\": wide_opt, \"deeptabular\": deep_opt}\n",
+    "schedulers = {\"wide\": wide_sch, \"deeptabular\": deep_sch}\n",
+    "initializers = {\"wide\": KaimingNormal, \"deeptabular\": XavierNormal}\n",
     "# General settings as List\n",
-    "callbacks = [LRHistory(n_epochs=10), EarlyStopping, ModelCheckpoint(filepath='model_weights/wd_out')]\n",
+    "callbacks = [\n",
+    "    LRHistory(n_epochs=10),\n",
+    "    EarlyStopping,\n",
+    "    ModelCheckpoint(filepath=\"model_weights/wd_out\"),\n",
+    "]\n",
     "metrics = [Accuracy, Recall]"
    ]
   },
@@ -641,14 +662,15 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "trainer = Trainer(model, \n",
-    "                  objective='binary', \n",
-    "                  optimizers=optimizers, \n",
-    "                  lr_schedulers=schedulers,\n",
-    "                  initializers=initializers,\n",
-    "                  callbacks=callbacks,\n",
-    "                  metrics=metrics\n",
-    "                 )"
+    "trainer = Trainer(\n",
+    "    model,\n",
+    "    objective=\"binary\",\n",
+    "    optimizers=optimizers,\n",
+    "    lr_schedulers=schedulers,\n",
+    "    initializers=initializers,\n",
+    "    callbacks=callbacks,\n",
+    "    metrics=metrics,\n",
+    ")"
    ]
   },
   {
@@ -691,7 +713,14 @@
     }
    ],
    "source": [
-    "trainer.fit(X_wide=X_wide, X_tab=X_tab, target=target, n_epochs=10, batch_size=256, val_split=0.2)"
+    "trainer.fit(\n",
+    "    X_wide=X_wide,\n",
+    "    X_tab=X_tab,\n",
+    "    target=target,\n",
+    "    n_epochs=10,\n",
+    "    batch_size=256,\n",
+    "    val_split=0.2,\n",
+    ")"
    ]
   },
   {
@@ -835,7 +864,9 @@
     }
    ],
    "source": [
-    "trainer.get_embeddings(col_name='education', cat_encoding_dict=tab_preprocessor.label_encoder.encoding_dict)"
+    "trainer.get_embeddings(\n",
+    "    col_name=\"education\", cat_encoding_dict=tab_preprocessor.label_encoder.encoding_dict\n",
+    ")"
    ]
   }
  ],
diff --git a/examples/05_Regression_with_Images_and_Text.ipynb b/examples/05_Regression_with_Images_and_Text.ipynb
index e6664b50..9b73806a 100644
--- a/examples/05_Regression_with_Images_and_Text.ipynb
+++ b/examples/05_Regression_with_Images_and_Text.ipynb
@@ -34,8 +34,20 @@
     "import torch\n",
     "\n",
     "from pytorch_widedeep import Trainer\n",
-    "from pytorch_widedeep.preprocessing import WidePreprocessor, TabPreprocessor, TextPreprocessor, ImagePreprocessor\n",
-    "from pytorch_widedeep.models import Wide, TabMlp, TabResnet, DeepText, DeepImage, WideDeep\n",
+    "from pytorch_widedeep.preprocessing import (\n",
+    "    WidePreprocessor,\n",
+    "    TabPreprocessor,\n",
+    "    TextPreprocessor,\n",
+    "    ImagePreprocessor,\n",
+    ")\n",
+    "from pytorch_widedeep.models import (\n",
+    "    Wide,\n",
+    "    TabMlp,\n",
+    "    TabResnet,\n",
+    "    DeepText,\n",
+    "    DeepImage,\n",
+    "    WideDeep,\n",
+    ")\n",
     "from pytorch_widedeep.losses import RMSELoss\n",
     "from pytorch_widedeep.initializers import *\n",
     "from pytorch_widedeep.callbacks import *\n",
@@ -952,7 +964,7 @@
     }
    ],
    "source": [
-    "df = pd.read_csv('data/airbnb/airbnb_sample.csv')\n",
+    "df = pd.read_csv(\"data/airbnb/airbnb_sample.csv\")\n",
     "df.head()"
    ]
   },
@@ -972,23 +984,30 @@
    "outputs": [],
    "source": [
     "# There are a number of columns that are already binary. Therefore, no need to one hot encode them\n",
-    "crossed_cols = [('property_type', 'room_type')]\n",
-    "already_dummies = [c for c in df.columns if 'amenity' in c] + ['has_house_rules']\n",
-    "wide_cols = ['is_location_exact', 'property_type', 'room_type', 'host_gender',\n",
-    "'instant_bookable'] + already_dummies\n",
-    "cat_embed_cols = [(c, 16) for c in df.columns if 'catg' in c] + \\\n",
-    "    [('neighbourhood_cleansed', 64), ('cancellation_policy', 16)]\n",
-    "continuous_cols = ['latitude', 'longitude', 'security_deposit', 'extra_people']\n",
+    "crossed_cols = [(\"property_type\", \"room_type\")]\n",
+    "already_dummies = [c for c in df.columns if \"amenity\" in c] + [\"has_house_rules\"]\n",
+    "wide_cols = [\n",
+    "    \"is_location_exact\",\n",
+    "    \"property_type\",\n",
+    "    \"room_type\",\n",
+    "    \"host_gender\",\n",
+    "    \"instant_bookable\",\n",
+    "] + already_dummies\n",
+    "cat_embed_cols = [(c, 16) for c in df.columns if \"catg\" in c] + [\n",
+    "    (\"neighbourhood_cleansed\", 64),\n",
+    "    (\"cancellation_policy\", 16),\n",
+    "]\n",
+    "continuous_cols = [\"latitude\", \"longitude\", \"security_deposit\", \"extra_people\"]\n",
     "# it does not make sense to standarised Latitude and Longitude\n",
-    "already_standard = ['latitude', 'longitude']\n",
-    "#text and image colnames\n",
-    "text_col = 'description'\n",
-    "img_col = 'id'\n",
+    "already_standard = [\"latitude\", \"longitude\"]\n",
+    "# text and image colnames\n",
+    "text_col = \"description\"\n",
+    "img_col = \"id\"\n",
     "# path to pretrained word embeddings and the images\n",
-    "word_vectors_path = 'data/glove.6B/glove.6B.100d.txt'\n",
-    "img_path = 'data/airbnb/property_picture'\n",
+    "word_vectors_path = \"data/glove.6B/glove.6B.100d.txt\"\n",
+    "img_path = \"data/airbnb/property_picture\"\n",
     "# target\n",
-    "target_col = 'yield'"
+    "target_col = \"yield\""
    ]
   },
   {
@@ -1027,7 +1046,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "tab_preprocessor = TabPreprocessor(embed_cols=cat_embed_cols, continuous_cols=continuous_cols)\n",
+    "tab_preprocessor = TabPreprocessor(\n",
+    "    embed_cols=cat_embed_cols, continuous_cols=continuous_cols\n",
+    ")\n",
     "X_tab = tab_preprocessor.fit_transform(df)"
    ]
   },
@@ -1049,7 +1070,9 @@
     }
    ],
    "source": [
-    "text_preprocessor = TextPreprocessor(word_vectors_path=word_vectors_path, text_col=text_col)\n",
+    "text_preprocessor = TextPreprocessor(\n",
+    "    word_vectors_path=word_vectors_path, text_col=text_col\n",
+    ")\n",
     "X_text = text_preprocessor.fit_transform(df)"
    ]
   },
@@ -1095,7 +1118,7 @@
     }
    ],
    "source": [
-    "image_processor = ImagePreprocessor(img_col = img_col, img_path = img_path)\n",
+    "image_processor = ImagePreprocessor(img_col=img_col, img_path=img_path)\n",
     "X_images = image_processor.fit_transform(df)"
    ]
   },
@@ -1117,23 +1140,24 @@
     "\n",
     "# DeepDense: 2 Dense layers\n",
     "deeptabular = TabMlp(\n",
-    "    column_idx = tab_preprocessor.column_idx,\n",
-    "    mlp_hidden_dims=[128,64],\n",
-    "    mlp_dropout = 0.1,\n",
-    "    mlp_batchnorm = True,\n",
+    "    column_idx=tab_preprocessor.column_idx,\n",
+    "    mlp_hidden_dims=[128, 64],\n",
+    "    mlp_dropout=0.1,\n",
+    "    mlp_batchnorm=True,\n",
     "    embed_input=tab_preprocessor.embeddings_input,\n",
-    "    embed_dropout = 0.1,\n",
-    "    continuous_cols = continuous_cols,\n",
-    "    cont_norm_layer = \"batchnorm\"\n",
+    "    embed_dropout=0.1,\n",
+    "    continuous_cols=continuous_cols,\n",
+    "    cont_norm_layer=\"batchnorm\",\n",
     ")\n",
-    "    \n",
+    "\n",
     "# DeepText: a stack of 2 LSTMs\n",
     "deeptext = DeepText(\n",
-    "    vocab_size=len(text_preprocessor.vocab.itos), \n",
-    "    hidden_dim=64, \n",
+    "    vocab_size=len(text_preprocessor.vocab.itos),\n",
+    "    hidden_dim=64,\n",
     "    n_layers=2,\n",
-    "    rnn_dropout=0.5, \n",
-    "    embed_matrix=text_preprocessor.embedding_matrix)\n",
+    "    rnn_dropout=0.5,\n",
+    "    embed_matrix=text_preprocessor.embedding_matrix,\n",
+    ")\n",
     "\n",
     "# Pretrained Resnet 18 (default is all but last 2 conv blocks frozen) plus a FC-Head 512->256->128\n",
     "deepimage = DeepImage(pretrained=True, head_hidden_dims=[512, 256, 128])"
@@ -1152,7 +1176,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "model = WideDeep(wide=wide, deeptabular=deeptabular, deeptext=deeptext, deepimage=deepimage)"
+    "model = WideDeep(\n",
+    "    wide=wide, deeptabular=deeptabular, deeptext=deeptext, deepimage=deepimage\n",
+    ")"
    ]
   },
   {
@@ -1168,7 +1194,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "trainer = Trainer(model, objective='rmse')"
+    "trainer = Trainer(model, objective=\"rmse\")"
    ]
   },
   {
@@ -1186,8 +1212,16 @@
     }
    ],
    "source": [
-    "trainer.fit(X_wide=X_wide, X_tab=X_tab, X_text=X_text, X_img=X_images,\n",
-    "    target=target, n_epochs=1, batch_size=32, val_split=0.2)"
+    "trainer.fit(\n",
+    "    X_wide=X_wide,\n",
+    "    X_tab=X_tab,\n",
+    "    X_text=X_text,\n",
+    "    X_img=X_images,\n",
+    "    target=target,\n",
+    "    n_epochs=1,\n",
+    "    batch_size=32,\n",
+    "    val_split=0.2,\n",
+    ")"
    ]
   },
   {
@@ -1208,23 +1242,23 @@
     "wide = Wide(wide_dim=np.unique(X_wide).shape[0], pred_dim=1)\n",
     "\n",
     "deeptabular = TabMlp(\n",
-    "    column_idx = tab_preprocessor.column_idx,\n",
-    "    mlp_hidden_dims=[128,64],\n",
-    "    mlp_dropout = 0.1,\n",
-    "    mlp_batchnorm = True,\n",
+    "    column_idx=tab_preprocessor.column_idx,\n",
+    "    mlp_hidden_dims=[128, 64],\n",
+    "    mlp_dropout=0.1,\n",
+    "    mlp_batchnorm=True,\n",
     "    embed_input=tab_preprocessor.embeddings_input,\n",
-    "    embed_dropout = 0.1,\n",
-    "    continuous_cols = continuous_cols,\n",
-    "    cont_norm_layer = \"batchnorm\"\n",
+    "    embed_dropout=0.1,\n",
+    "    continuous_cols=continuous_cols,\n",
+    "    cont_norm_layer=\"batchnorm\",\n",
     ")\n",
     "\n",
     "deeptext = DeepText(\n",
-    "    vocab_size=len(text_preprocessor.vocab.itos), \n",
-    "    hidden_dim=128, \n",
-    "    n_layers=2, \n",
-    "    rnn_dropout=0.5, \n",
-    "    embed_matrix=text_preprocessor.embedding_matrix\n",
-    "    )\n",
+    "    vocab_size=len(text_preprocessor.vocab.itos),\n",
+    "    hidden_dim=128,\n",
+    "    n_layers=2,\n",
+    "    rnn_dropout=0.5,\n",
+    "    embed_matrix=text_preprocessor.embedding_matrix,\n",
+    ")\n",
     "\n",
     "deepimage = DeepImage(pretrained=True, head_hidden_dims=[512, 256, 128])"
    ]
@@ -1243,11 +1277,11 @@
    "outputs": [],
    "source": [
     "model = WideDeep(\n",
-    "    wide=wide, \n",
-    "    deeptabular=deeptabular, \n",
-    "    deeptext=deeptext, \n",
-    "    deepimage=deepimage, \n",
-    "    head_hidden_dims=[128, 64]\n",
+    "    wide=wide,\n",
+    "    deeptabular=deeptabular,\n",
+    "    deeptext=deeptext,\n",
+    "    deepimage=deepimage,\n",
+    "    head_hidden_dims=[128, 64],\n",
     ")"
    ]
   },
@@ -1451,10 +1485,12 @@
    "source": [
     "deep_params = []\n",
     "for childname, child in model.named_children():\n",
-    "    if childname == 'deeptabular':\n",
-    "        for n,p in child.named_parameters():\n",
-    "            if \"embed_layer\" in n: deep_params.append({'params': p, 'lr': 1e-4})\n",
-    "            else: deep_params.append({'params': p, 'lr': 1e-3})"
+    "    if childname == \"deeptabular\":\n",
+    "        for n, p in child.named_parameters():\n",
+    "            if \"embed_layer\" in n:\n",
+    "                deep_params.append({\"params\": p, \"lr\": 1e-4})\n",
+    "            else:\n",
+    "                deep_params.append({\"params\": p, \"lr\": 1e-3})"
    ]
   },
   {
@@ -1466,7 +1502,7 @@
     "wide_opt = torch.optim.Adam(model.wide.parameters(), lr=0.03)\n",
     "deep_opt = torch.optim.Adam(deep_params)\n",
     "text_opt = RAdam(model.deeptext.parameters())\n",
-    "img_opt  = RAdam(model.deepimage.parameters())\n",
+    "img_opt = RAdam(model.deepimage.parameters())\n",
     "head_opt = torch.optim.Adam(model.deephead.parameters())"
    ]
   },
@@ -1477,9 +1513,9 @@
    "outputs": [],
    "source": [
     "wide_sch = torch.optim.lr_scheduler.StepLR(wide_opt, step_size=5)\n",
-    "deep_sch = torch.optim.lr_scheduler.MultiStepLR(deep_opt, milestones=[3,8])\n",
+    "deep_sch = torch.optim.lr_scheduler.MultiStepLR(deep_opt, milestones=[3, 8])\n",
     "text_sch = torch.optim.lr_scheduler.StepLR(text_opt, step_size=5)\n",
-    "img_sch  = torch.optim.lr_scheduler.MultiStepLR(deep_opt, milestones=[3,8])\n",
+    "img_sch = torch.optim.lr_scheduler.MultiStepLR(deep_opt, milestones=[3, 8])\n",
     "head_sch = torch.optim.lr_scheduler.StepLR(head_opt, step_size=5)"
    ]
   },
@@ -1490,23 +1526,42 @@
    "outputs": [],
    "source": [
     "# remember, one optimizer per model components, for lr_schedures and initializers is not neccesary\n",
-    "optimizers = {'wide': wide_opt, 'deeptabular':deep_opt, 'deeptext':text_opt, 'deepimage': img_opt, 'deephead': head_opt}\n",
-    "schedulers = {'wide': wide_sch, 'deeptabular':deep_sch, 'deeptext':text_sch, 'deepimage': img_sch, 'deephead': head_sch}\n",
+    "optimizers = {\n",
+    "    \"wide\": wide_opt,\n",
+    "    \"deeptabular\": deep_opt,\n",
+    "    \"deeptext\": text_opt,\n",
+    "    \"deepimage\": img_opt,\n",
+    "    \"deephead\": head_opt,\n",
+    "}\n",
+    "schedulers = {\n",
+    "    \"wide\": wide_sch,\n",
+    "    \"deeptabular\": deep_sch,\n",
+    "    \"deeptext\": text_sch,\n",
+    "    \"deepimage\": img_sch,\n",
+    "    \"deephead\": head_sch,\n",
+    "}\n",
     "\n",
     "# Now...we have used pretrained word embeddings, so you do not want to\n",
     "# initialise these  embeddings. However you might still want to initialise the\n",
     "# other layers in the DeepText component. No probs, you can do that with the\n",
     "# parameter pattern and your knowledge on regular  expressions. Here we are\n",
     "# telling to the KaimingNormal initializer to NOT touch the  parameters whose\n",
-    "# name contains the string word_embed. \n",
-    "initializers = {'wide': KaimingNormal, 'deeptabular':KaimingNormal, \n",
-    "                'deeptext':KaimingNormal(pattern=r\"^(?!.*word_embed).*$\"), \n",
-    "                'deepimage':KaimingNormal}\n",
+    "# name contains the string word_embed.\n",
+    "initializers = {\n",
+    "    \"wide\": KaimingNormal,\n",
+    "    \"deeptabular\": KaimingNormal,\n",
+    "    \"deeptext\": KaimingNormal(pattern=r\"^(?!.*word_embed).*$\"),\n",
+    "    \"deepimage\": KaimingNormal,\n",
+    "}\n",
     "\n",
-    "mean = [0.406, 0.456, 0.485]  #BGR\n",
-    "std =  [0.225, 0.224, 0.229]  #BGR\n",
+    "mean = [0.406, 0.456, 0.485]  # BGR\n",
+    "std = [0.225, 0.224, 0.229]  # BGR\n",
     "transforms = [ToTensor, Normalize(mean=mean, std=std)]\n",
-    "callbacks = [LRHistory(n_epochs=10), EarlyStopping, ModelCheckpoint(filepath='model_weights/wd_out')]"
+    "callbacks = [\n",
+    "    LRHistory(n_epochs=10),\n",
+    "    EarlyStopping,\n",
+    "    ModelCheckpoint(filepath=\"model_weights/wd_out\"),\n",
+    "]"
    ]
   },
   {
@@ -1525,13 +1580,13 @@
    ],
    "source": [
     "trainer = Trainer(\n",
-    "    model, \n",
-    "    objective=\"rmse\", \n",
-    "    initializers=initializers,    \n",
-    "    optimizers=optimizers,    \n",
+    "    model,\n",
+    "    objective=\"rmse\",\n",
+    "    initializers=initializers,\n",
+    "    optimizers=optimizers,\n",
     "    lr_schedulers=schedulers,\n",
-    "    callbacks=callbacks, \n",
-    "    transforms=transforms\n",
+    "    callbacks=callbacks,\n",
+    "    transforms=transforms,\n",
     ")"
    ]
   },
@@ -1564,8 +1619,16 @@
     }
    ],
    "source": [
-    "trainer.fit(X_wide=X_wide, X_tab=X_tab, X_text=X_text, X_img=X_images,\n",
-    "    target=target, n_epochs=1, batch_size=32, val_split=0.2)"
+    "trainer.fit(\n",
+    "    X_wide=X_wide,\n",
+    "    X_tab=X_tab,\n",
+    "    X_text=X_text,\n",
+    "    X_img=X_images,\n",
+    "    target=target,\n",
+    "    n_epochs=1,\n",
+    "    batch_size=32,\n",
+    "    val_split=0.2,\n",
+    ")"
    ]
   },
   {
diff --git a/examples/06_FineTune_and_WarmUp_Model_Components.ipynb b/examples/06_FineTune_and_WarmUp_Model_Components.ipynb
index 1395e388..db442928 100644
--- a/examples/06_FineTune_and_WarmUp_Model_Components.ipynb
+++ b/examples/06_FineTune_and_WarmUp_Model_Components.ipynb
@@ -226,12 +226,12 @@
     }
    ],
    "source": [
-    "df = pd.read_csv('data/adult/adult.csv.zip')\n",
+    "df = pd.read_csv(\"data/adult/adult.csv.zip\")\n",
     "# For convenience, we'll replace '-' with '_'\n",
     "df.columns = [c.replace(\"-\", \"_\") for c in df.columns]\n",
-    "#binary target\n",
-    "df['income_label'] = (df[\"income\"].apply(lambda x: \">50K\" in x)).astype(int)\n",
-    "df.drop('income', axis=1, inplace=True)\n",
+    "# binary target\n",
+    "df[\"income_label\"] = (df[\"income\"].apply(lambda x: \">50K\" in x)).astype(int)\n",
+    "df.drop(\"income\", axis=1, inplace=True)\n",
     "df.head()"
    ]
   },
@@ -241,11 +241,24 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "wide_cols = ['education', 'relationship','workclass','occupation','native_country','gender']\n",
-    "crossed_cols = [('education', 'occupation'), ('native_country', 'occupation')]\n",
-    "cat_embed_cols = [('education',16), ('relationship',8), ('workclass',16), ('occupation',16),('native_country',16)]\n",
-    "continuous_cols = [\"age\",\"hours_per_week\"]\n",
-    "target_col = 'income_label'"
+    "wide_cols = [\n",
+    "    \"education\",\n",
+    "    \"relationship\",\n",
+    "    \"workclass\",\n",
+    "    \"occupation\",\n",
+    "    \"native_country\",\n",
+    "    \"gender\",\n",
+    "]\n",
+    "crossed_cols = [(\"education\", \"occupation\"), (\"native_country\", \"occupation\")]\n",
+    "cat_embed_cols = [\n",
+    "    (\"education\", 16),\n",
+    "    (\"relationship\", 8),\n",
+    "    (\"workclass\", 16),\n",
+    "    (\"occupation\", 16),\n",
+    "    (\"native_country\", 16),\n",
+    "]\n",
+    "continuous_cols = [\"age\", \"hours_per_week\"]\n",
+    "target_col = \"income_label\""
    ]
   },
   {
@@ -262,7 +275,9 @@
     "X_wide = wide_preprocessor.fit_transform(df)\n",
     "\n",
     "# DEEP\n",
-    "tab_preprocessor = TabPreprocessor(embed_cols=cat_embed_cols, continuous_cols=continuous_cols)\n",
+    "tab_preprocessor = TabPreprocessor(\n",
+    "    embed_cols=cat_embed_cols, continuous_cols=continuous_cols\n",
+    ")\n",
     "X_tab = tab_preprocessor.fit_transform(df)"
    ]
   },
@@ -273,11 +288,12 @@
    "outputs": [],
    "source": [
     "wide = Wide(wide_dim=np.unique(X_wide).shape[0], pred_dim=1)\n",
-    "deeptabular = TabMlp(mlp_hidden_dims=[64,32], \n",
-    "                   column_idx=tab_preprocessor.column_idx,\n",
-    "                   embed_input=tab_preprocessor.embeddings_input,\n",
-    "                   continuous_cols=continuous_cols\n",
-    "                    )\n",
+    "deeptabular = TabMlp(\n",
+    "    mlp_hidden_dims=[64, 32],\n",
+    "    column_idx=tab_preprocessor.column_idx,\n",
+    "    embed_input=tab_preprocessor.embeddings_input,\n",
+    "    continuous_cols=continuous_cols,\n",
+    ")\n",
     "model = WideDeep(wide=wide, deeptabular=deeptabular)"
    ]
   },
@@ -314,7 +330,9 @@
     }
    ],
    "source": [
-    "trainer.fit(X_wide=X_wide, X_tab=X_tab, target=target, n_epochs=2, val_split=0.2, batch_size=256)"
+    "trainer.fit(\n",
+    "    X_wide=X_wide, X_tab=X_tab, target=target, n_epochs=2, val_split=0.2, batch_size=256\n",
+    ")"
    ]
   },
   {
@@ -342,10 +360,12 @@
    "outputs": [],
    "source": [
     "wide_1 = Wide(wide_dim=np.unique(X_wide).shape[0], pred_dim=1)\n",
-    "deeptabular_1 = TabMlp(mlp_hidden_dims=[64,32], \n",
-    "                   column_idx=tab_preprocessor.column_idx,\n",
-    "                   embed_input=tab_preprocessor.embeddings_input,\n",
-    "                   continuous_cols=continuous_cols)\n",
+    "deeptabular_1 = TabMlp(\n",
+    "    mlp_hidden_dims=[64, 32],\n",
+    "    column_idx=tab_preprocessor.column_idx,\n",
+    "    embed_input=tab_preprocessor.embeddings_input,\n",
+    "    continuous_cols=continuous_cols,\n",
+    ")\n",
     "model_1 = WideDeep(wide=wide_1, deeptabular=deeptabular_1)"
    ]
   },
@@ -440,13 +460,14 @@
    ],
    "source": [
     "trainer_1.fit(\n",
-    "    X_wide=X_wide, \n",
-    "    X_tab=X_tab, \n",
-    "    target=target, \n",
-    "    finetune=True, \n",
-    "    finetune_epochs=2, \n",
-    "    n_epochs=2, \n",
-    "    batch_size=256)"
+    "    X_wide=X_wide,\n",
+    "    X_tab=X_tab,\n",
+    "    target=target,\n",
+    "    finetune=True,\n",
+    "    finetune_epochs=2,\n",
+    "    n_epochs=2,\n",
+    "    batch_size=256,\n",
+    ")"
    ]
   },
   {
@@ -463,10 +484,12 @@
    "outputs": [],
    "source": [
     "wide = Wide(wide_dim=np.unique(X_wide).shape[0], pred_dim=1)\n",
-    "deeptabular = TabMlp(mlp_hidden_dims=[128, 32], \n",
-    "                   column_idx=tab_preprocessor.column_idx,\n",
-    "                   embed_input=tab_preprocessor.embeddings_input,\n",
-    "                   continuous_cols=continuous_cols)\n",
+    "deeptabular = TabMlp(\n",
+    "    mlp_hidden_dims=[128, 32],\n",
+    "    column_idx=tab_preprocessor.column_idx,\n",
+    "    embed_input=tab_preprocessor.embeddings_input,\n",
+    "    continuous_cols=continuous_cols,\n",
+    ")\n",
     "model = WideDeep(wide=wide, deeptabular=deeptabular)"
    ]
   },
@@ -543,14 +566,14 @@
    ],
    "source": [
     "trainer_2.fit(\n",
-    "    X_wide=X_wide, \n",
-    "    X_tab=X_tab, \n",
-    "    target=target, \n",
-    "    val_split=0.1, \n",
-    "    warmup=True, \n",
-    "    warmup_epochs=2, \n",
-    "    n_epochs=2, \n",
-    "    batch_size=256\n",
+    "    X_wide=X_wide,\n",
+    "    X_tab=X_tab,\n",
+    "    target=target,\n",
+    "    val_split=0.1,\n",
+    "    warmup=True,\n",
+    "    warmup_epochs=2,\n",
+    "    n_epochs=2,\n",
+    "    batch_size=256,\n",
     ")"
    ]
   },
@@ -608,10 +631,11 @@
    "source": [
     "wide = Wide(wide_dim=np.unique(X_wide).shape[0], pred_dim=1)\n",
     "deeptabular = TabResnet(\n",
-    "    blocks_dims=[128, 64, 32], \n",
+    "    blocks_dims=[128, 64, 32],\n",
     "    column_idx=tab_preprocessor.column_idx,\n",
     "    embed_input=tab_preprocessor.embeddings_input,\n",
-    "    continuous_cols=continuous_cols)\n",
+    "    continuous_cols=continuous_cols,\n",
+    ")\n",
     "model = WideDeep(wide=wide, deeptabular=deeptabular)"
    ]
   },
@@ -718,7 +742,9 @@
     }
    ],
    "source": [
-    "trainer_3.fit(X_wide=X_wide, X_tab=X_tab, target=target, val_split=0.1, n_epochs=2, batch_size=256)"
+    "trainer_3.fit(\n",
+    "    X_wide=X_wide, X_tab=X_tab, target=target, val_split=0.1, n_epochs=2, batch_size=256\n",
+    ")"
    ]
   },
   {
@@ -747,10 +773,11 @@
    "source": [
     "wide_3 = Wide(wide_dim=np.unique(X_wide).shape[0], pred_dim=1)\n",
     "deeptabular_3 = TabResnet(\n",
-    "    blocks_dims=[128, 64, 32], \n",
+    "    blocks_dims=[128, 64, 32],\n",
     "    column_idx=tab_preprocessor.column_idx,\n",
     "    embed_input=tab_preprocessor.embeddings_input,\n",
-    "    continuous_cols=continuous_cols)\n",
+    "    continuous_cols=continuous_cols,\n",
+    ")\n",
     "model_3 = WideDeep(wide=wide, deeptabular=deeptabular)"
    ]
   },
@@ -958,17 +985,17 @@
    ],
    "source": [
     "trainer_4.fit(\n",
-    "    X_wide=X_wide, \n",
-    "    X_tab=X_tab, \n",
-    "    target=target, \n",
-    "    val_split=0.1, \n",
-    "    finetune=True, \n",
-    "    finetune_epochs=2, \n",
+    "    X_wide=X_wide,\n",
+    "    X_tab=X_tab,\n",
+    "    target=target,\n",
+    "    val_split=0.1,\n",
+    "    finetune=True,\n",
+    "    finetune_epochs=2,\n",
     "    finetune_deeptabular_gradual=True,\n",
-    "    finetune_deeptabular_layers = tab_layers,\n",
-    "    finetune_deeptabular_max_lr = 0.01,\n",
+    "    finetune_deeptabular_layers=tab_layers,\n",
+    "    finetune_deeptabular_max_lr=0.01,\n",
     "    n_epochs=2,\n",
-    "    batch_size=256\n",
+    "    batch_size=256,\n",
     ")"
    ]
   },
@@ -985,10 +1012,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "deeptabular = TabMlp(mlp_hidden_dims=[200, 100], \n",
-    "                   column_idx=tab_preprocessor.column_idx,\n",
-    "                   embed_input=tab_preprocessor.embeddings_input,\n",
-    "                   continuous_cols=continuous_cols)\n",
+    "deeptabular = TabMlp(\n",
+    "    mlp_hidden_dims=[200, 100],\n",
+    "    column_idx=tab_preprocessor.column_idx,\n",
+    "    embed_input=tab_preprocessor.embeddings_input,\n",
+    "    continuous_cols=continuous_cols,\n",
+    ")\n",
     "model = WideDeep(deeptabular=deeptabular)"
    ]
   },
@@ -1016,7 +1045,9 @@
     }
    ],
    "source": [
-    "trainer_5.fit(X_wide=X_wide, X_tab=X_tab, target=target, val_split=0.1, n_epochs=1, batch_size=256)"
+    "trainer_5.fit(\n",
+    "    X_wide=X_wide, X_tab=X_tab, target=target, val_split=0.1, n_epochs=1, batch_size=256\n",
+    ")"
    ]
   },
   {
@@ -1034,10 +1065,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "deeptabular_5 = TabMlp(mlp_hidden_dims=[200, 100], \n",
-    "                   column_idx=tab_preprocessor.column_idx,\n",
-    "                   embed_input=tab_preprocessor.embeddings_input,\n",
-    "                   continuous_cols=continuous_cols)\n",
+    "deeptabular_5 = TabMlp(\n",
+    "    mlp_hidden_dims=[200, 100],\n",
+    "    column_idx=tab_preprocessor.column_idx,\n",
+    "    embed_input=tab_preprocessor.embeddings_input,\n",
+    "    continuous_cols=continuous_cols,\n",
+    ")\n",
     "model_5 = WideDeep(deeptabular=deeptabular_5)"
    ]
   },
@@ -1121,17 +1154,16 @@
    ],
    "source": [
     "trainer_6.fit(\n",
-    "    X_wide=X_wide, \n",
-    "    X_tab=X_tab, \n",
-    "    target=target, \n",
-    "    val_split=0.1, \n",
-    "    finetune=True, \n",
+    "    X_wide=X_wide,\n",
+    "    X_tab=X_tab,\n",
+    "    target=target,\n",
+    "    val_split=0.1,\n",
+    "    finetune=True,\n",
     "    finetune_epochs=2,\n",
     "    finetune_max_lr=0.01,\n",
     "    stop_after_finetuning=True,\n",
-    "    batch_size=256\n",
-    "    \n",
-    ")   "
+    "    batch_size=256,\n",
+    ")"
    ]
   },
   {
diff --git a/examples/07_Custom_Components.ipynb b/examples/07_Custom_Components.ipynb
index eb71c4cb..5b490981 100644
--- a/examples/07_Custom_Components.ipynb
+++ b/examples/07_Custom_Components.ipynb
@@ -34,8 +34,20 @@
     "import torch\n",
     "\n",
     "from pytorch_widedeep import Trainer\n",
-    "from pytorch_widedeep.preprocessing import WidePreprocessor, TabPreprocessor, TextPreprocessor, ImagePreprocessor\n",
-    "from pytorch_widedeep.models import Wide, TabMlp, TabResnet, DeepText, DeepImage, WideDeep\n",
+    "from pytorch_widedeep.preprocessing import (\n",
+    "    WidePreprocessor,\n",
+    "    TabPreprocessor,\n",
+    "    TextPreprocessor,\n",
+    "    ImagePreprocessor,\n",
+    ")\n",
+    "from pytorch_widedeep.models import (\n",
+    "    Wide,\n",
+    "    TabMlp,\n",
+    "    TabResnet,\n",
+    "    DeepText,\n",
+    "    DeepImage,\n",
+    "    WideDeep,\n",
+    ")\n",
     "from pytorch_widedeep.losses import RMSELoss\n",
     "from pytorch_widedeep.initializers import *\n",
     "from pytorch_widedeep.callbacks import *"
@@ -951,7 +963,7 @@
     }
    ],
    "source": [
-    "df = pd.read_csv('data/airbnb/airbnb_sample.csv')\n",
+    "df = pd.read_csv(\"data/airbnb/airbnb_sample.csv\")\n",
     "df.head()"
    ]
   },
@@ -962,23 +974,30 @@
    "outputs": [],
    "source": [
     "# There are a number of columns that are already binary. Therefore, no need to one hot encode them\n",
-    "crossed_cols = [('property_type', 'room_type')]\n",
-    "already_dummies = [c for c in df.columns if 'amenity' in c] + ['has_house_rules']\n",
-    "wide_cols = ['is_location_exact', 'property_type', 'room_type', 'host_gender',\n",
-    "'instant_bookable'] + already_dummies\n",
-    "cat_embed_cols = [(c, 16) for c in df.columns if 'catg' in c] + \\\n",
-    "    [('neighbourhood_cleansed', 64), ('cancellation_policy', 16)]\n",
-    "continuous_cols = ['latitude', 'longitude', 'security_deposit', 'extra_people']\n",
+    "crossed_cols = [(\"property_type\", \"room_type\")]\n",
+    "already_dummies = [c for c in df.columns if \"amenity\" in c] + [\"has_house_rules\"]\n",
+    "wide_cols = [\n",
+    "    \"is_location_exact\",\n",
+    "    \"property_type\",\n",
+    "    \"room_type\",\n",
+    "    \"host_gender\",\n",
+    "    \"instant_bookable\",\n",
+    "] + already_dummies\n",
+    "cat_embed_cols = [(c, 16) for c in df.columns if \"catg\" in c] + [\n",
+    "    (\"neighbourhood_cleansed\", 64),\n",
+    "    (\"cancellation_policy\", 16),\n",
+    "]\n",
+    "continuous_cols = [\"latitude\", \"longitude\", \"security_deposit\", \"extra_people\"]\n",
     "# it does not make sense to standarised Latitude and Longitude\n",
-    "already_standard = ['latitude', 'longitude']\n",
-    "# text and image colnames\n",
-    "text_col = 'description'\n",
-    "img_col = 'id'\n",
+    "already_standard = [\"latitude\", \"longitude\"]\n",
+    "#  text and image colnames\n",
+    "text_col = \"description\"\n",
+    "img_col = \"id\"\n",
     "# path to pretrained word embeddings and the images\n",
-    "word_vectors_path = 'data/glove.6B/glove.6B.100d.txt'\n",
-    "img_path = 'data/airbnb/property_picture'\n",
+    "word_vectors_path = \"data/glove.6B/glove.6B.100d.txt\"\n",
+    "img_path = \"data/airbnb/property_picture\"\n",
     "# target\n",
-    "target_col = 'yield'"
+    "target_col = \"yield\""
    ]
   },
   {
@@ -1040,13 +1059,17 @@
     "wide_preprocessor = WidePreprocessor(wide_cols=wide_cols, crossed_cols=crossed_cols)\n",
     "X_wide = wide_preprocessor.fit_transform(df)\n",
     "\n",
-    "tab_preprocessor = TabPreprocessor(embed_cols=cat_embed_cols, continuous_cols=continuous_cols)\n",
+    "tab_preprocessor = TabPreprocessor(\n",
+    "    embed_cols=cat_embed_cols, continuous_cols=continuous_cols\n",
+    ")\n",
     "X_tab = tab_preprocessor.fit_transform(df)\n",
     "\n",
-    "text_preprocessor = TextPreprocessor(word_vectors_path=word_vectors_path, text_col=text_col)\n",
+    "text_preprocessor = TextPreprocessor(\n",
+    "    word_vectors_path=word_vectors_path, text_col=text_col\n",
+    ")\n",
     "X_text = text_preprocessor.fit_transform(df)\n",
     "\n",
-    "image_processor = ImagePreprocessor(img_col = img_col, img_path = img_path)\n",
+    "image_processor = ImagePreprocessor(img_col=img_col, img_path=img_path)\n",
     "X_images = image_processor.fit_transform(df)"
    ]
   },
@@ -1070,16 +1093,16 @@
     "\n",
     "# DeepDense: 2 Dense layers\n",
     "deeptabular = TabMlp(\n",
-    "    column_idx = tab_preprocessor.column_idx,\n",
-    "    mlp_hidden_dims=[128,64],\n",
-    "    mlp_dropout = 0.1,\n",
-    "    mlp_batchnorm = True,\n",
+    "    column_idx=tab_preprocessor.column_idx,\n",
+    "    mlp_hidden_dims=[128, 64],\n",
+    "    mlp_dropout=0.1,\n",
+    "    mlp_batchnorm=True,\n",
     "    embed_input=tab_preprocessor.embeddings_input,\n",
-    "    embed_dropout = 0.1,\n",
-    "    continuous_cols = continuous_cols,\n",
-    "    cont_norm_layer = \"batchnorm\"\n",
+    "    embed_dropout=0.1,\n",
+    "    continuous_cols=continuous_cols,\n",
+    "    cont_norm_layer=\"batchnorm\",\n",
     ")\n",
-    "   \n",
+    "\n",
     "# Pretrained Resnet 18 (default is all but last 2 conv blocks frozen) plus a FC-Head 512->256->128\n",
     "deepimage = DeepImage(pretrained=True, head_hidden_dims=[512, 256, 128])"
    ]
@@ -1104,9 +1127,7 @@
     "        super(MyDeepText, self).__init__()\n",
     "\n",
     "        # word/token embeddings\n",
-    "        self.word_embed = nn.Embedding(\n",
-    "            vocab_size, embed_dim, padding_idx=padding_idx\n",
-    "        )\n",
+    "        self.word_embed = nn.Embedding(vocab_size, embed_dim, padding_idx=padding_idx)\n",
     "\n",
     "        # stack of RNNs\n",
     "        self.rnn = nn.GRU(\n",
@@ -1141,7 +1162,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "model = WideDeep(wide=wide, deeptabular=deeptabular, deeptext=mydeeptext, deepimage=deepimage)"
+    "model = WideDeep(\n",
+    "    wide=wide, deeptabular=deeptabular, deeptext=mydeeptext, deepimage=deepimage\n",
+    ")"
    ]
   },
   {
@@ -1182,7 +1205,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "trainer = Trainer(model, objective='regression', custom_loss_function=RMSELoss())"
+    "trainer = Trainer(model, objective=\"regression\", custom_loss_function=RMSELoss())"
    ]
   },
   {
@@ -1200,8 +1223,16 @@
     }
    ],
    "source": [
-    "trainer.fit(X_wide=X_wide, X_tab=X_tab, X_text=X_text, X_img=X_images,\n",
-    "    target=target, n_epochs=1, batch_size=32, val_split=0.2)"
+    "trainer.fit(\n",
+    "    X_wide=X_wide,\n",
+    "    X_tab=X_tab,\n",
+    "    X_text=X_text,\n",
+    "    X_img=X_images,\n",
+    "    target=target,\n",
+    "    n_epochs=1,\n",
+    "    batch_size=32,\n",
+    "    val_split=0.2,\n",
+    ")"
    ]
   },
   {
@@ -1389,7 +1420,7 @@
     }
    ],
    "source": [
-    "df = pd.read_csv('data/adult/adult.csv.zip')\n",
+    "df = pd.read_csv(\"data/adult/adult.csv.zip\")\n",
     "df.head()"
    ]
   },
@@ -1562,9 +1593,9 @@
    "source": [
     "# For convenience, we'll replace '-' with '_'\n",
     "df.columns = [c.replace(\"-\", \"_\") for c in df.columns]\n",
-    "# binary target\n",
-    "df['income_label'] = (df[\"income\"].apply(lambda x: \">50K\" in x)).astype(int)\n",
-    "df.drop('income', axis=1, inplace=True)\n",
+    "#  binary target\n",
+    "df[\"income_label\"] = (df[\"income\"].apply(lambda x: \">50K\" in x)).astype(int)\n",
+    "df.drop(\"income\", axis=1, inplace=True)\n",
     "df.head()"
    ]
   },
@@ -1574,11 +1605,24 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "wide_cols = ['education', 'relationship','workclass','occupation','native_country','gender']\n",
-    "crossed_cols = [('education', 'occupation'), ('native_country', 'occupation')]\n",
-    "cat_embed_cols = [('education',16), ('relationship',8), ('workclass',16), ('occupation',16),('native_country',16)]\n",
-    "continuous_cols = [\"age\",\"hours_per_week\"]\n",
-    "target_col = 'income_label'"
+    "wide_cols = [\n",
+    "    \"education\",\n",
+    "    \"relationship\",\n",
+    "    \"workclass\",\n",
+    "    \"occupation\",\n",
+    "    \"native_country\",\n",
+    "    \"gender\",\n",
+    "]\n",
+    "crossed_cols = [(\"education\", \"occupation\"), (\"native_country\", \"occupation\")]\n",
+    "cat_embed_cols = [\n",
+    "    (\"education\", 16),\n",
+    "    (\"relationship\", 8),\n",
+    "    (\"workclass\", 16),\n",
+    "    (\"occupation\", 16),\n",
+    "    (\"native_country\", 16),\n",
+    "]\n",
+    "continuous_cols = [\"age\", \"hours_per_week\"]\n",
+    "target_col = \"income_label\""
    ]
   },
   {
@@ -1595,7 +1639,9 @@
     "X_wide = wide_preprocessor.fit_transform(df)\n",
     "\n",
     "# deeptabular\n",
-    "tab_preprocessor = TabPreprocessor(embed_cols=cat_embed_cols, continuous_cols=continuous_cols)\n",
+    "tab_preprocessor = TabPreprocessor(\n",
+    "    embed_cols=cat_embed_cols, continuous_cols=continuous_cols\n",
+    ")\n",
     "X_tab = tab_preprocessor.fit_transform(df)"
    ]
   },
@@ -1606,10 +1652,12 @@
    "outputs": [],
    "source": [
     "wide = Wide(wide_dim=np.unique(X_wide).shape[0], pred_dim=1)\n",
-    "deeptabular = TabMlp(mlp_hidden_dims=[64,32], \n",
-    "                      column_idx=tab_preprocessor.column_idx,\n",
-    "                      embed_input=tab_preprocessor.embeddings_input,\n",
-    "                      continuous_cols=continuous_cols)\n",
+    "deeptabular = TabMlp(\n",
+    "    mlp_hidden_dims=[64, 32],\n",
+    "    column_idx=tab_preprocessor.column_idx,\n",
+    "    embed_input=tab_preprocessor.embeddings_input,\n",
+    "    continuous_cols=continuous_cols,\n",
+    ")\n",
     "model = WideDeep(wide=wide, deeptabular=deeptabular)"
    ]
   },
@@ -1667,7 +1715,7 @@
     "        self.correct_count += y_pred.eq(y_true).sum().item()\n",
     "        self.total_count += len(y_pred)\n",
     "        accuracy = float(self.correct_count) / float(self.total_count)\n",
-    "        return np.array(accuracy)\n"
+    "        return np.array(accuracy)"
    ]
   },
   {
@@ -1696,18 +1744,18 @@
    "outputs": [],
    "source": [
     "class SillyCallback(Callback):\n",
-    "    def on_train_begin(self, logs = None):\n",
+    "    def on_train_begin(self, logs=None):\n",
     "        # recordings will be the trainer object attributes\n",
     "        self.trainer.silly_callback = {}\n",
     "\n",
-    "        self.trainer.silly_callback['beginning'] = []\n",
-    "        self.trainer.silly_callback['end'] = []\n",
+    "        self.trainer.silly_callback[\"beginning\"] = []\n",
+    "        self.trainer.silly_callback[\"end\"] = []\n",
     "\n",
     "    def on_epoch_begin(self, epoch, logs=None):\n",
-    "        self.trainer.silly_callback['beginning'].append(epoch+1)\n",
+    "        self.trainer.silly_callback[\"beginning\"].append(epoch + 1)\n",
     "\n",
     "    def on_epoch_end(self, epoch, logs=None, metric=None):\n",
-    "        self.trainer.silly_callback['end'].append(epoch+1)"
+    "        self.trainer.silly_callback[\"end\"].append(epoch + 1)"
    ]
   },
   {
@@ -1723,7 +1771,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "trainer = Trainer(model, objective='binary', metrics=[Accuracy], callbacks=[SillyCallback])"
+    "trainer = Trainer(\n",
+    "    model, objective=\"binary\", metrics=[Accuracy], callbacks=[SillyCallback]\n",
+    ")"
    ]
   },
   {
@@ -1749,7 +1799,9 @@
     }
    ],
    "source": [
-    "trainer.fit(X_wide=X_wide, X_tab=X_tab, target=target, n_epochs=5, batch_size=64, val_split=0.2)"
+    "trainer.fit(\n",
+    "    X_wide=X_wide, X_tab=X_tab, target=target, n_epochs=5, batch_size=64, val_split=0.2\n",
+    ")"
    ]
   },
   {
diff --git a/examples/08_save_and_load_model_and_artifacts.ipynb b/examples/08_save_and_load_model_and_artifacts.ipynb
index 2b4605e4..dc3a2c29 100644
--- a/examples/08_save_and_load_model_and_artifacts.ipynb
+++ b/examples/08_save_and_load_model_and_artifacts.ipynb
@@ -205,7 +205,7 @@
     }
    ],
    "source": [
-    "df = pd.read_csv('data/adult/adult.csv.zip')\n",
+    "df = pd.read_csv(\"data/adult/adult.csv.zip\")\n",
     "df.head()"
    ]
   },
@@ -378,9 +378,9 @@
    "source": [
     "# For convenience, we'll replace '-' with '_'\n",
     "df.columns = [c.replace(\"-\", \"_\") for c in df.columns]\n",
-    "#binary target\n",
-    "df['target'] = (df[\"income\"].apply(lambda x: \">50K\" in x)).astype(int)\n",
-    "df.drop('income', axis=1, inplace=True)\n",
+    "# binary target\n",
+    "df[\"target\"] = (df[\"income\"].apply(lambda x: \">50K\" in x)).astype(int)\n",
+    "df.drop(\"income\", axis=1, inplace=True)\n",
     "df.head()"
    ]
   },
@@ -422,8 +422,15 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "wide_cols = ['education', 'relationship','workclass','occupation','native_country','gender']\n",
-    "crossed_cols = [('education', 'occupation'), ('native_country', 'occupation')]"
+    "wide_cols = [\n",
+    "    \"education\",\n",
+    "    \"relationship\",\n",
+    "    \"workclass\",\n",
+    "    \"occupation\",\n",
+    "    \"native_country\",\n",
+    "    \"gender\",\n",
+    "]\n",
+    "crossed_cols = [(\"education\", \"occupation\"), (\"native_country\", \"occupation\")]"
    ]
   },
   {
@@ -456,7 +463,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "tab_preprocessor = TabPreprocessor(embed_cols=cat_embed_cols, continuous_cols=num_cols, scale=True)\n",
+    "tab_preprocessor = TabPreprocessor(\n",
+    "    embed_cols=cat_embed_cols, continuous_cols=num_cols, scale=True\n",
+    ")\n",
     "X_tab_train = tab_preprocessor.fit_transform(train)\n",
     "y_train = train.target.values\n",
     "X_tab_valid = tab_preprocessor.transform(valid)\n",
@@ -668,7 +677,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "trainer.save(path=\"tmp_dir/\", model_filename=\"model_state_dict_saved_option_2.pt\", save_state_dict=True)"
+    "trainer.save(\n",
+    "    path=\"tmp_dir/\",\n",
+    "    model_filename=\"model_state_dict_saved_option_2.pt\",\n",
+    "    save_state_dict=True,\n",
+    ")"
    ]
   },
   {
@@ -736,7 +749,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "with open('tmp_dir/wide_preproc.pkl', 'wb') as wp:\n",
+    "with open(\"tmp_dir/wide_preproc.pkl\", \"wb\") as wp:\n",
     "    pickle.dump(wide_preprocessor, wp)"
    ]
   },
@@ -746,7 +759,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "with open('tmp_dir/tab_preproc.pkl', 'wb') as dp:\n",
+    "with open(\"tmp_dir/tab_preproc.pkl\", \"wb\") as dp:\n",
     "    pickle.dump(tab_preprocessor, dp)"
    ]
   },
@@ -756,7 +769,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "with open('tmp_dir/eary_stop.pkl', 'wb') as es:\n",
+    "with open(\"tmp_dir/eary_stop.pkl\", \"wb\") as es:\n",
     "    pickle.dump(early_stopping, es)"
    ]
   },
@@ -977,7 +990,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "with open('tmp_dir/wide_preproc.pkl', 'rb') as wp:\n",
+    "with open(\"tmp_dir/wide_preproc.pkl\", \"rb\") as wp:\n",
     "    wide_preprocessor_new = pickle.load(wp)"
    ]
   },
@@ -987,7 +1000,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "with open('tmp_dir/tab_preproc.pkl', 'rb') as tp:\n",
+    "with open(\"tmp_dir/tab_preproc.pkl\", \"rb\") as tp:\n",
     "    tab_preprocessor_new = pickle.load(tp)"
    ]
   },
diff --git a/examples/09_Custom_DataLoader_Imbalanced_dataset.ipynb b/examples/09_Custom_DataLoader_Imbalanced_dataset.ipynb
index 1203e0b9..23e36c13 100644
--- a/examples/09_Custom_DataLoader_Imbalanced_dataset.ipynb
+++ b/examples/09_Custom_DataLoader_Imbalanced_dataset.ipynb
@@ -66,11 +66,12 @@
     "import datetime\n",
     "\n",
     "import warnings\n",
+    "\n",
     "warnings.filterwarnings(\"ignore\", category=DeprecationWarning)\n",
     "\n",
     "# increase displayed columns in jupyter notebook\n",
-    "pd.set_option('display.max_columns', 200)\n",
-    "pd.set_option('display.max_rows', 300)"
+    "pd.set_option(\"display.max_columns\", 200)\n",
+    "pd.set_option(\"display.max_rows\", 300)"
    ]
   },
   {
@@ -640,8 +641,8 @@
     }
    ],
    "source": [
-    "header_list = ['EXAMPLE_ID', 'BLOCK_ID', 'target'] + [str(i) for i in range(4,78)]\n",
-    "df = pd.read_csv('data/kddcup04/bio_train.dat', sep='\\t', names=header_list)\n",
+    "header_list = [\"EXAMPLE_ID\", \"BLOCK_ID\", \"target\"] + [str(i) for i in range(4, 78)]\n",
+    "df = pd.read_csv(\"data/kddcup04/bio_train.dat\", sep=\"\\t\", names=header_list)\n",
     "df.head()"
    ]
   },
@@ -665,7 +666,7 @@
    ],
    "source": [
     "# imbalance of the classes\n",
-    "df['target'].value_counts()"
+    "df[\"target\"].value_counts()"
    ]
   },
   {
@@ -675,7 +676,7 @@
    "outputs": [],
    "source": [
     "# drop columns we won't need in this example\n",
-    "df.drop(columns=['EXAMPLE_ID', 'BLOCK_ID'], inplace=True)"
+    "df.drop(columns=[\"EXAMPLE_ID\", \"BLOCK_ID\"], inplace=True)"
    ]
   },
   {
@@ -684,8 +685,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "df_train, df_valid = train_test_split(df, test_size=0.2, stratify=df['target'], random_state=1)\n",
-    "df_valid, df_test = train_test_split(df_valid, test_size=0.5, stratify=df_valid['target'], random_state=1)"
+    "df_train, df_valid = train_test_split(\n",
+    "    df, test_size=0.2, stratify=df[\"target\"], random_state=1\n",
+    ")\n",
+    "df_valid, df_test = train_test_split(\n",
+    "    df_valid, test_size=0.5, stratify=df_valid[\"target\"], random_state=1\n",
+    ")"
    ]
   },
   {
@@ -701,7 +706,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "continuous_cols = df.drop(columns=['target']).columns.values.tolist()"
+    "continuous_cols = df.drop(columns=[\"target\"]).columns.values.tolist()"
    ]
   },
   {
@@ -711,16 +716,15 @@
    "outputs": [],
    "source": [
     "# deeptabular\n",
-    "tab_preprocessor = TabPreprocessor(continuous_cols=continuous_cols,\n",
-    "                                   scale=True)\n",
+    "tab_preprocessor = TabPreprocessor(continuous_cols=continuous_cols, scale=True)\n",
     "X_tab_train = tab_preprocessor.fit_transform(df_train)\n",
     "X_tab_valid = tab_preprocessor.transform(df_valid)\n",
     "X_tab_test = tab_preprocessor.transform(df_test)\n",
     "\n",
     "# target\n",
-    "y_train = df_train['target'].values\n",
-    "y_valid = df_valid['target'].values\n",
-    "y_test = df_test['target'].values"
+    "y_train = df_train[\"target\"].values\n",
+    "y_valid = df_valid[\"target\"].values\n",
+    "y_test = df_test[\"target\"].values"
    ]
   },
   {
@@ -738,7 +742,9 @@
    "source": [
     "input_layer = len(tab_preprocessor.continuous_cols)\n",
     "output_layer = 1\n",
-    "hidden_layers = np.linspace(input_layer*2, output_layer, 5, endpoint=False, dtype=int).tolist()"
+    "hidden_layers = np.linspace(\n",
+    "    input_layer * 2, output_layer, 5, endpoint=False, dtype=int\n",
+    ").tolist()"
    ]
   },
   {
@@ -796,9 +802,11 @@
     }
    ],
    "source": [
-    "deeptabular = TabMlp(mlp_hidden_dims=hidden_layers,\n",
-    "                     column_idx=tab_preprocessor.column_idx,\n",
-    "                     continuous_cols=tab_preprocessor.continuous_cols)\n",
+    "deeptabular = TabMlp(\n",
+    "    mlp_hidden_dims=hidden_layers,\n",
+    "    column_idx=tab_preprocessor.column_idx,\n",
+    "    continuous_cols=tab_preprocessor.continuous_cols,\n",
+    ")\n",
     "model = WideDeep(deeptabular=deeptabular)\n",
     "model"
    ]
@@ -811,7 +819,7 @@
    "source": [
     "# Metrics from torchmetrics\n",
     "accuracy = Accuracy_torchmetrics(average=None, num_classes=2)\n",
-    "precision = Precision_torchmetrics(average='micro', num_classes=2)\n",
+    "precision = Precision_torchmetrics(average=\"micro\", num_classes=2)\n",
     "f1 = F1_torchmetrics(average=None, num_classes=2)\n",
     "recall = Recall_torchmetrics(average=None, num_classes=2)"
    ]
@@ -840,13 +848,15 @@
     "# LR Scheduler\n",
     "deep_sch = lr_scheduler.StepLR(deep_opt, step_size=3)\n",
     "\n",
-    "trainer = Trainer(model,\n",
-    "                  objective=\"binary\",\n",
-    "                  lr_schedulers={'deeptabular':deep_sch},\n",
-    "                  initializers={'deeptabular':XavierNormal},\n",
-    "                  optimizers={'deeptabular':deep_opt},\n",
-    "                  metrics=[accuracy, precision, recall, f1],\n",
-    "                  verbose=1)"
+    "trainer = Trainer(\n",
+    "    model,\n",
+    "    objective=\"binary\",\n",
+    "    lr_schedulers={\"deeptabular\": deep_sch},\n",
+    "    initializers={\"deeptabular\": XavierNormal},\n",
+    "    optimizers={\"deeptabular\": deep_opt},\n",
+    "    metrics=[accuracy, precision, recall, f1],\n",
+    "    verbose=1,\n",
+    ")"
    ]
   },
   {
@@ -881,14 +891,19 @@
    ],
    "source": [
     "start = time.time()\n",
-    "trainer.fit(X_train={\"X_tab\": X_tab_train, \"target\": y_train},\n",
-    "            X_val={\"X_tab\": X_tab_valid, \"target\": y_valid},\n",
-    "            n_epochs=3,\n",
-    "            batch_size=50,\n",
-    "            custom_dataloader=DataLoaderImbalanced,\n",
-    "            oversample_mul=5\n",
-    "           )\n",
-    "print('Training time[s]: {}'.format(datetime.timedelta(seconds=round(time.time()-start))))"
+    "trainer.fit(\n",
+    "    X_train={\"X_tab\": X_tab_train, \"target\": y_train},\n",
+    "    X_val={\"X_tab\": X_tab_valid, \"target\": y_valid},\n",
+    "    n_epochs=3,\n",
+    "    batch_size=50,\n",
+    "    custom_dataloader=DataLoaderImbalanced,\n",
+    "    oversample_mul=5,\n",
+    ")\n",
+    "print(\n",
+    "    \"Training time[s]: {}\".format(\n",
+    "        datetime.timedelta(seconds=round(time.time() - start))\n",
+    "    )\n",
+    ")"
    ]
   },
   {
@@ -1041,7 +1056,7 @@
    ],
    "source": [
     "df_pred = trainer.predict(X_tab=X_tab_test)\n",
-    "print(classification_report(df_test['target'].to_list(), df_pred))\n",
+    "print(classification_report(df_test[\"target\"].to_list(), df_pred))\n",
     "print(\"Actual predicted values:\\n{}\".format(np.unique(df_pred, return_counts=True)))"
    ]
   }
diff --git a/examples/10_The_Transformer_Family.ipynb b/examples/10_The_Transformer_Family.ipynb
index c6f922b9..122f9c50 100644
--- a/examples/10_The_Transformer_Family.ipynb
+++ b/examples/10_The_Transformer_Family.ipynb
@@ -34,7 +34,14 @@
     "\n",
     "from pytorch_widedeep.preprocessing import TabPreprocessor\n",
     "from pytorch_widedeep.training import Trainer\n",
-    "from pytorch_widedeep.models import TabTransformer, SAINT, FTTransformer, TabFastFormer, TabPerceiver, WideDeep\n",
+    "from pytorch_widedeep.models import (\n",
+    "    TabTransformer,\n",
+    "    SAINT,\n",
+    "    FTTransformer,\n",
+    "    TabFastFormer,\n",
+    "    TabPerceiver,\n",
+    "    WideDeep,\n",
+    ")\n",
     "from pytorch_widedeep.metrics import Accuracy"
    ]
   },
@@ -205,7 +212,7 @@
     }
    ],
    "source": [
-    "df = pd.read_csv('data/adult/adult.csv.zip')\n",
+    "df = pd.read_csv(\"data/adult/adult.csv.zip\")\n",
     "df.head()"
    ]
   },
@@ -372,8 +379,8 @@
    "source": [
     "# For convenience, we'll replace '-' with '_'\n",
     "df.columns = [c.replace(\"-\", \"_\") for c in df.columns]\n",
-    "#binary target\n",
-    "df['target'] = (df[\"income\"].apply(lambda x: \">50K\" in x)).astype(int)\n",
+    "# binary target\n",
+    "df[\"target\"] = (df[\"income\"].apply(lambda x: \">50K\" in x)).astype(int)\n",
     "df.drop([\"income\", \"educational_num\"], axis=1, inplace=True)\n",
     "\n",
     "df.head()"
@@ -390,7 +397,7 @@
     "    # 50 is just a random number I choose here for this example\n",
     "    if df[col].dtype == \"O\" or df[col].nunique() < 50 and col != \"target\":\n",
     "        cat_cols.append(col)\n",
-    "    elif col != \"target\":        \n",
+    "    elif col != \"target\":\n",
     "        cont_cols.append(col)\n",
     "target_col = \"target\""
    ]
@@ -410,10 +417,9 @@
    "source": [
     "target = df[target_col].values\n",
     "\n",
-    "tab_preprocessor = TabPreprocessor(embed_cols=cat_cols, \n",
-    "                                   continuous_cols=cont_cols, \n",
-    "                                   for_transformer=True\n",
-    "                                  )\n",
+    "tab_preprocessor = TabPreprocessor(\n",
+    "    embed_cols=cat_cols, continuous_cols=cont_cols, for_transformer=True\n",
+    ")\n",
     "X_tab = tab_preprocessor.fit_transform(df)"
    ]
   },
@@ -423,14 +429,16 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# here all categorical columns will be encoded as 32 dim embeddings, then passed through the transformer \n",
+    "# here all categorical columns will be encoded as 32 dim embeddings, then passed through the transformer\n",
     "# blocks, concatenated with the continuous and finally through an MLP\n",
-    "tab_transformer = TabTransformer(column_idx=tab_preprocessor.column_idx,\n",
-    "                                 embed_input=tab_preprocessor.embeddings_input,\n",
-    "                                 continuous_cols=tab_preprocessor.continuous_cols, \n",
-    "                                 cont_norm_layer=\"batchnorm\", \n",
-    "                                 n_blocks=4, n_heads=4 \n",
-    "                                )"
+    "tab_transformer = TabTransformer(\n",
+    "    column_idx=tab_preprocessor.column_idx,\n",
+    "    embed_input=tab_preprocessor.embeddings_input,\n",
+    "    continuous_cols=tab_preprocessor.continuous_cols,\n",
+    "    cont_norm_layer=\"batchnorm\",\n",
+    "    n_blocks=4,\n",
+    "    n_heads=4,\n",
+    ")"
    ]
   },
   {
@@ -582,7 +590,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "trainer = Trainer(model, objective='binary', metrics=[Accuracy])"
+    "trainer = Trainer(model, objective=\"binary\", metrics=[Accuracy])"
    ]
   },
   {
@@ -616,11 +624,12 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "tab_preprocessor = TabPreprocessor(embed_cols=cat_cols, \n",
-    "                                   continuous_cols=cont_cols, \n",
-    "                                   for_transformer=True,\n",
-    "                                   with_cls_token=True\n",
-    "                                  )\n",
+    "tab_preprocessor = TabPreprocessor(\n",
+    "    embed_cols=cat_cols,\n",
+    "    continuous_cols=cont_cols,\n",
+    "    for_transformer=True,\n",
+    "    with_cls_token=True,\n",
+    ")\n",
     "X_tab = tab_preprocessor.fit_transform(df)"
    ]
   },
@@ -630,11 +639,14 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "ft_transformer = FTTransformer(column_idx=tab_preprocessor.column_idx,\n",
-    "                                embed_input=tab_preprocessor.embeddings_input,\n",
-    "                                continuous_cols=tab_preprocessor.continuous_cols, \n",
-    "                                n_blocks=3, n_heads=6, input_dim=36\n",
-    "                               )"
+    "ft_transformer = FTTransformer(\n",
+    "    column_idx=tab_preprocessor.column_idx,\n",
+    "    embed_input=tab_preprocessor.embeddings_input,\n",
+    "    continuous_cols=tab_preprocessor.continuous_cols,\n",
+    "    n_blocks=3,\n",
+    "    n_heads=6,\n",
+    "    input_dim=36,\n",
+    ")"
    ]
   },
   {
@@ -748,7 +760,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "trainer = Trainer(model, objective='binary', metrics=[Accuracy])"
+    "trainer = Trainer(model, objective=\"binary\", metrics=[Accuracy])"
    ]
   },
   {
@@ -782,12 +794,14 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "saint = SAINT(column_idx=tab_preprocessor.column_idx,\n",
-    "              embed_input=tab_preprocessor.embeddings_input,\n",
-    "              continuous_cols=tab_preprocessor.continuous_cols, \n",
-    "              transformer_activation=\"geglu\",\n",
-    "              n_blocks=2, n_heads=4, \n",
-    "             )"
+    "saint = SAINT(\n",
+    "    column_idx=tab_preprocessor.column_idx,\n",
+    "    embed_input=tab_preprocessor.embeddings_input,\n",
+    "    continuous_cols=tab_preprocessor.continuous_cols,\n",
+    "    transformer_activation=\"geglu\",\n",
+    "    n_blocks=2,\n",
+    "    n_heads=4,\n",
+    ")"
    ]
   },
   {
@@ -937,7 +951,7 @@
    ],
    "source": [
     "model = WideDeep(deeptabular=saint)\n",
-    "trainer = Trainer(model, objective='binary', metrics=[Accuracy])\n",
+    "trainer = Trainer(model, objective=\"binary\", metrics=[Accuracy])\n",
     "trainer.fit(X_tab=X_tab, target=target, n_epochs=1, batch_size=128, val_split=0.2)"
    ]
   },
@@ -958,11 +972,13 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "tabfastformer = TabFastFormer(column_idx=tab_preprocessor.column_idx,\n",
-    "              embed_input=tab_preprocessor.embeddings_input,\n",
-    "              continuous_cols=tab_preprocessor.continuous_cols, \n",
-    "              n_blocks=2, n_heads=4,             \n",
-    "            )"
+    "tabfastformer = TabFastFormer(\n",
+    "    column_idx=tab_preprocessor.column_idx,\n",
+    "    embed_input=tab_preprocessor.embeddings_input,\n",
+    "    continuous_cols=tab_preprocessor.continuous_cols,\n",
+    "    n_blocks=2,\n",
+    "    n_heads=4,\n",
+    ")"
    ]
   },
   {
@@ -1078,7 +1094,7 @@
    ],
    "source": [
     "model = WideDeep(deeptabular=tabfastformer)\n",
-    "trainer = Trainer(model, objective='binary', metrics=[Accuracy])\n",
+    "trainer = Trainer(model, objective=\"binary\", metrics=[Accuracy])\n",
     "trainer.fit(X_tab=X_tab, target=target, n_epochs=1, batch_size=256, val_split=0.2)"
    ]
   },
@@ -1097,10 +1113,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "tab_preprocessor = TabPreprocessor(embed_cols=cat_cols, \n",
-    "                                   continuous_cols=cont_cols, \n",
-    "                                   for_transformer=True,\n",
-    "                                  )\n",
+    "tab_preprocessor = TabPreprocessor(\n",
+    "    embed_cols=cat_cols,\n",
+    "    continuous_cols=cont_cols,\n",
+    "    for_transformer=True,\n",
+    ")\n",
     "X_tab = tab_preprocessor.fit_transform(df)"
    ]
   },
@@ -1113,10 +1130,10 @@
     "tabperceiver = TabPerceiver(\n",
     "    column_idx=tab_preprocessor.column_idx,\n",
     "    embed_input=tab_preprocessor.embeddings_input,\n",
-    "    continuous_cols=tab_preprocessor.continuous_cols, \n",
-    "    n_perceiver_blocks=1, \n",
-    "    n_latent_blocks=3, \n",
-    "    n_latent_heads=2, \n",
+    "    continuous_cols=tab_preprocessor.continuous_cols,\n",
+    "    n_perceiver_blocks=1,\n",
+    "    n_latent_blocks=3,\n",
+    "    n_latent_heads=2,\n",
     "    n_latents=6,\n",
     "    latent_dim=32,\n",
     ")"
@@ -1138,7 +1155,7 @@
    ],
    "source": [
     "model = WideDeep(deeptabular=tabperceiver)\n",
-    "trainer = Trainer(model, objective='binary', metrics=[Accuracy])\n",
+    "trainer = Trainer(model, objective=\"binary\", metrics=[Accuracy])\n",
     "trainer.fit(X_tab=X_tab, target=target, n_epochs=1, batch_size=256, val_split=0.2)"
    ]
   },
@@ -1157,12 +1174,13 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "tab_preprocessor = TabPreprocessor(embed_cols=cat_cols, \n",
-    "                                   continuous_cols=cont_cols, \n",
-    "                                   for_transformer=True,\n",
-    "                                   shared_embed=True,\n",
-    "                                   with_cls_token=True\n",
-    "                                  )\n",
+    "tab_preprocessor = TabPreprocessor(\n",
+    "    embed_cols=cat_cols,\n",
+    "    continuous_cols=cont_cols,\n",
+    "    for_transformer=True,\n",
+    "    shared_embed=True,\n",
+    "    with_cls_token=True,\n",
+    ")\n",
     "X_tab = tab_preprocessor.fit_transform(df)"
    ]
   },
@@ -1172,15 +1190,17 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "ft_transformer = TabTransformer(column_idx=tab_preprocessor.column_idx,\n",
-    "                                embed_input=tab_preprocessor.embeddings_input,\n",
-    "                                continuous_cols=tab_preprocessor.continuous_cols, \n",
-    "                                embed_continuous=True,\n",
-    "                                embed_continuous_activation=None,       \n",
-    "                                shared_embed=True,  \n",
-    "                                cont_norm_layer=\"batchnorm\", \n",
-    "                                n_blocks=4, n_heads=4 \n",
-    "                                )"
+    "ft_transformer = TabTransformer(\n",
+    "    column_idx=tab_preprocessor.column_idx,\n",
+    "    embed_input=tab_preprocessor.embeddings_input,\n",
+    "    continuous_cols=tab_preprocessor.continuous_cols,\n",
+    "    embed_continuous=True,\n",
+    "    embed_continuous_activation=None,\n",
+    "    shared_embed=True,\n",
+    "    cont_norm_layer=\"batchnorm\",\n",
+    "    n_blocks=4,\n",
+    "    n_heads=4,\n",
+    ")"
    ]
   },
   {
@@ -1370,7 +1390,7 @@
    ],
    "source": [
     "model = WideDeep(deeptabular=ft_transformer)\n",
-    "trainer = Trainer(model, objective='binary', metrics=[Accuracy])\n",
+    "trainer = Trainer(model, objective=\"binary\", metrics=[Accuracy])\n",
     "trainer.fit(X_tab=X_tab, target=target, n_epochs=1, batch_size=256, val_split=0.2)"
    ]
   }
diff --git a/examples/11_Extracting_Embeddings.ipynb b/examples/11_Extracting_Embeddings.ipynb
index 5504f84d..b0097ea5 100644
--- a/examples/11_Extracting_Embeddings.ipynb
+++ b/examples/11_Extracting_Embeddings.ipynb
@@ -193,7 +193,7 @@
     }
    ],
    "source": [
-    "df = pd.read_csv('data/adult/adult.csv.zip')\n",
+    "df = pd.read_csv(\"data/adult/adult.csv.zip\")\n",
     "df.head()"
    ]
   },
@@ -360,8 +360,8 @@
    "source": [
     "# For convenience, we'll replace '-' with '_'\n",
     "df.columns = [c.replace(\"-\", \"_\") for c in df.columns]\n",
-    "#binary target\n",
-    "df['target'] = (df[\"income\"].apply(lambda x: \">50K\" in x)).astype(int)\n",
+    "# binary target\n",
+    "df[\"target\"] = (df[\"income\"].apply(lambda x: \">50K\" in x)).astype(int)\n",
     "df.drop([\"income\", \"educational_num\"], axis=1, inplace=True)\n",
     "\n",
     "df.head()"
@@ -378,7 +378,7 @@
     "    # 50 is just a random number I choose here for this example\n",
     "    if df[col].dtype == \"O\" or df[col].nunique() < 50 and col != \"target\":\n",
     "        cat_cols.append(col)\n",
-    "    elif col != \"target\":        \n",
+    "    elif col != \"target\":\n",
     "        cont_cols.append(col)\n",
     "target_col = \"target\""
    ]
@@ -391,10 +391,9 @@
    "source": [
     "target = df[target_col].values\n",
     "\n",
-    "tab_preprocessor = TabPreprocessor(embed_cols=cat_cols, \n",
-    "                                   continuous_cols=cont_cols, \n",
-    "                                   for_transformer=True\n",
-    "                                  )\n",
+    "tab_preprocessor = TabPreprocessor(\n",
+    "    embed_cols=cat_cols, continuous_cols=cont_cols, for_transformer=True\n",
+    ")\n",
     "X_tab = tab_preprocessor.fit_transform(df)"
    ]
   },
@@ -404,11 +403,14 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "ft_transformer = FTTransformer(column_idx=tab_preprocessor.column_idx,\n",
-    "                                embed_input=tab_preprocessor.embeddings_input,\n",
-    "                                continuous_cols=tab_preprocessor.continuous_cols, \n",
-    "                                n_blocks=3, n_heads=6, input_dim=36\n",
-    "                               )"
+    "ft_transformer = FTTransformer(\n",
+    "    column_idx=tab_preprocessor.column_idx,\n",
+    "    embed_input=tab_preprocessor.embeddings_input,\n",
+    "    continuous_cols=tab_preprocessor.continuous_cols,\n",
+    "    n_blocks=3,\n",
+    "    n_heads=6,\n",
+    "    input_dim=36,\n",
+    ")"
    ]
   },
   {
@@ -427,7 +429,7 @@
    ],
    "source": [
     "model = WideDeep(deeptabular=ft_transformer)\n",
-    "trainer = Trainer(model, objective='binary', metrics=[Accuracy])\n",
+    "trainer = Trainer(model, objective=\"binary\", metrics=[Accuracy])\n",
     "trainer.fit(X_tab=X_tab, target=target, n_epochs=1, batch_size=256, val_split=0.2)"
    ]
   },
diff --git a/examples/12_HyperParameter_tuning_w_RayTune.ipynb b/examples/12_HyperParameter_tuning_w_RayTune.ipynb
new file mode 100644
index 00000000..618c52d4
--- /dev/null
+++ b/examples/12_HyperParameter_tuning_w_RayTune.ipynb
@@ -0,0 +1,1070 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Hyperparameter tuning and using Raytune and visulization using Tensorboard"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "* In this notebook we will use the higly imbalanced Protein Homology Dataset from [KDD cup 2004](https://www.kdd.org/kdd-cup/view/kdd-cup-2004/Data)\n",
+    "\n",
+    "```\n",
+    "* The first element of each line is a BLOCK ID that denotes to which native sequence this example belongs. There is a unique BLOCK ID for each native sequence. BLOCK IDs are integers running from 1 to 303 (one for each native sequence, i.e. for each query). BLOCK IDs were assigned before the blocks were split into the train and test sets, so they do not run consecutively in either file.\n",
+    "* The second element of each line is an EXAMPLE ID that uniquely describes the example. You will need this EXAMPLE ID and the BLOCK ID when you submit results.\n",
+    "* The third element is the class of the example. Proteins that are homologous to the native sequence are denoted by 1, non-homologous proteins (i.e. decoys) by 0. Test examples have a \"?\" in this position.\n",
+    "* All following elements are feature values. There are 74 feature values in each line. The features describe the match (e.g. the score of a sequence alignment) between the native protein sequence and the sequence that is tested for homology.\n",
+    "```"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Initial imports"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "/Users/javier/.pyenv/versions/3.7.7/lib/python3.7/importlib/_bootstrap.py:219: RuntimeWarning: numpy.ufunc size changed, may indicate binary incompatibility. Expected 192 from C header, got 216 from PyObject\n",
+      "  return f(*args, **kwds)\n"
+     ]
+    }
+   ],
+   "source": [
+    "import numpy as np\n",
+    "import pandas as pd\n",
+    "import torch\n",
+    "from torch.optim import SGD, lr_scheduler\n",
+    "\n",
+    "from pytorch_widedeep import Trainer\n",
+    "from pytorch_widedeep.preprocessing import TabPreprocessor\n",
+    "from pytorch_widedeep.models import TabMlp, WideDeep\n",
+    "from pytorch_widedeep.dataloaders import DataLoaderImbalanced, DataLoaderDefault\n",
+    "from torchmetrics import F1 as F1_torchmetrics\n",
+    "from torchmetrics import Accuracy as Accuracy_torchmetrics\n",
+    "from torchmetrics import Precision as Precision_torchmetrics\n",
+    "from torchmetrics import Recall as Recall_torchmetrics\n",
+    "from pytorch_widedeep.metrics import Accuracy, Recall, Precision, F1Score, R2Score\n",
+    "from pytorch_widedeep.initializers import XavierNormal\n",
+    "from pytorch_widedeep.callbacks import (\n",
+    "    EarlyStopping,\n",
+    "    ModelCheckpoint,\n",
+    "    LRHistory,\n",
+    "    RayTuneReporter,\n",
+    ")\n",
+    "\n",
+    "from sklearn.model_selection import train_test_split\n",
+    "from sklearn.metrics import classification_report\n",
+    "\n",
+    "import time\n",
+    "import datetime\n",
+    "\n",
+    "import warnings\n",
+    "\n",
+    "warnings.filterwarnings(\"ignore\", category=DeprecationWarning)\n",
+    "\n",
+    "from ray import tune\n",
+    "from ray.tune import JupyterNotebookReporter\n",
+    "import tracemalloc\n",
+    "\n",
+    "tracemalloc.start()\n",
+    "\n",
+    "# increase displayed columns in jupyter notebook\n",
+    "pd.set_option(\"display.max_columns\", 200)\n",
+    "pd.set_option(\"display.max_rows\", 300)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "<div>\n",
+       "<style scoped>\n",
+       "    .dataframe tbody tr th:only-of-type {\n",
+       "        vertical-align: middle;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe tbody tr th {\n",
+       "        vertical-align: top;\n",
+       "    }\n",
+       "\n",
+       "    .dataframe thead th {\n",
+       "        text-align: right;\n",
+       "    }\n",
+       "</style>\n",
+       "<table border=\"1\" class=\"dataframe\">\n",
+       "  <thead>\n",
+       "    <tr style=\"text-align: right;\">\n",
+       "      <th></th>\n",
+       "      <th>EXAMPLE_ID</th>\n",
+       "      <th>BLOCK_ID</th>\n",
+       "      <th>target</th>\n",
+       "      <th>4</th>\n",
+       "      <th>5</th>\n",
+       "      <th>6</th>\n",
+       "      <th>7</th>\n",
+       "      <th>8</th>\n",
+       "      <th>9</th>\n",
+       "      <th>10</th>\n",
+       "      <th>11</th>\n",
+       "      <th>12</th>\n",
+       "      <th>13</th>\n",
+       "      <th>14</th>\n",
+       "      <th>15</th>\n",
+       "      <th>16</th>\n",
+       "      <th>17</th>\n",
+       "      <th>18</th>\n",
+       "      <th>19</th>\n",
+       "      <th>20</th>\n",
+       "      <th>21</th>\n",
+       "      <th>22</th>\n",
+       "      <th>23</th>\n",
+       "      <th>24</th>\n",
+       "      <th>25</th>\n",
+       "      <th>26</th>\n",
+       "      <th>27</th>\n",
+       "      <th>28</th>\n",
+       "      <th>29</th>\n",
+       "      <th>30</th>\n",
+       "      <th>31</th>\n",
+       "      <th>32</th>\n",
+       "      <th>33</th>\n",
+       "      <th>34</th>\n",
+       "      <th>35</th>\n",
+       "      <th>36</th>\n",
+       "      <th>37</th>\n",
+       "      <th>38</th>\n",
+       "      <th>39</th>\n",
+       "      <th>40</th>\n",
+       "      <th>41</th>\n",
+       "      <th>42</th>\n",
+       "      <th>43</th>\n",
+       "      <th>44</th>\n",
+       "      <th>45</th>\n",
+       "      <th>46</th>\n",
+       "      <th>47</th>\n",
+       "      <th>48</th>\n",
+       "      <th>49</th>\n",
+       "      <th>50</th>\n",
+       "      <th>51</th>\n",
+       "      <th>52</th>\n",
+       "      <th>53</th>\n",
+       "      <th>54</th>\n",
+       "      <th>55</th>\n",
+       "      <th>56</th>\n",
+       "      <th>57</th>\n",
+       "      <th>58</th>\n",
+       "      <th>59</th>\n",
+       "      <th>60</th>\n",
+       "      <th>61</th>\n",
+       "      <th>62</th>\n",
+       "      <th>63</th>\n",
+       "      <th>64</th>\n",
+       "      <th>65</th>\n",
+       "      <th>66</th>\n",
+       "      <th>67</th>\n",
+       "      <th>68</th>\n",
+       "      <th>69</th>\n",
+       "      <th>70</th>\n",
+       "      <th>71</th>\n",
+       "      <th>72</th>\n",
+       "      <th>73</th>\n",
+       "      <th>74</th>\n",
+       "      <th>75</th>\n",
+       "      <th>76</th>\n",
+       "      <th>77</th>\n",
+       "    </tr>\n",
+       "  </thead>\n",
+       "  <tbody>\n",
+       "    <tr>\n",
+       "      <th>0</th>\n",
+       "      <td>279</td>\n",
+       "      <td>261532</td>\n",
+       "      <td>0</td>\n",
+       "      <td>52.0</td>\n",
+       "      <td>32.69</td>\n",
+       "      <td>0.30</td>\n",
+       "      <td>2.5</td>\n",
+       "      <td>20.0</td>\n",
+       "      <td>1256.8</td>\n",
+       "      <td>-0.89</td>\n",
+       "      <td>0.33</td>\n",
+       "      <td>11.0</td>\n",
+       "      <td>-55.0</td>\n",
+       "      <td>267.2</td>\n",
+       "      <td>0.52</td>\n",
+       "      <td>0.05</td>\n",
+       "      <td>-2.36</td>\n",
+       "      <td>49.6</td>\n",
+       "      <td>252.0</td>\n",
+       "      <td>0.43</td>\n",
+       "      <td>1.16</td>\n",
+       "      <td>-2.06</td>\n",
+       "      <td>-33.0</td>\n",
+       "      <td>-123.2</td>\n",
+       "      <td>1.60</td>\n",
+       "      <td>-0.49</td>\n",
+       "      <td>-6.06</td>\n",
+       "      <td>65.0</td>\n",
+       "      <td>296.1</td>\n",
+       "      <td>-0.28</td>\n",
+       "      <td>-0.26</td>\n",
+       "      <td>-3.83</td>\n",
+       "      <td>-22.6</td>\n",
+       "      <td>-170.0</td>\n",
+       "      <td>3.06</td>\n",
+       "      <td>-1.05</td>\n",
+       "      <td>-3.29</td>\n",
+       "      <td>22.9</td>\n",
+       "      <td>286.3</td>\n",
+       "      <td>0.12</td>\n",
+       "      <td>2.58</td>\n",
+       "      <td>4.08</td>\n",
+       "      <td>-33.0</td>\n",
+       "      <td>-178.9</td>\n",
+       "      <td>1.88</td>\n",
+       "      <td>0.53</td>\n",
+       "      <td>-7.0</td>\n",
+       "      <td>-44.0</td>\n",
+       "      <td>1987.0</td>\n",
+       "      <td>-5.41</td>\n",
+       "      <td>0.95</td>\n",
+       "      <td>-4.0</td>\n",
+       "      <td>-57.0</td>\n",
+       "      <td>722.9</td>\n",
+       "      <td>-3.26</td>\n",
+       "      <td>-0.55</td>\n",
+       "      <td>-7.5</td>\n",
+       "      <td>125.5</td>\n",
+       "      <td>1547.2</td>\n",
+       "      <td>-0.36</td>\n",
+       "      <td>1.12</td>\n",
+       "      <td>9.0</td>\n",
+       "      <td>-37.0</td>\n",
+       "      <td>72.5</td>\n",
+       "      <td>0.47</td>\n",
+       "      <td>0.74</td>\n",
+       "      <td>-11.0</td>\n",
+       "      <td>-8.0</td>\n",
+       "      <td>1595.1</td>\n",
+       "      <td>-1.64</td>\n",
+       "      <td>2.83</td>\n",
+       "      <td>-2.0</td>\n",
+       "      <td>-50.0</td>\n",
+       "      <td>445.2</td>\n",
+       "      <td>-0.35</td>\n",
+       "      <td>0.26</td>\n",
+       "      <td>0.76</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>1</th>\n",
+       "      <td>279</td>\n",
+       "      <td>261533</td>\n",
+       "      <td>0</td>\n",
+       "      <td>58.0</td>\n",
+       "      <td>33.33</td>\n",
+       "      <td>0.00</td>\n",
+       "      <td>16.5</td>\n",
+       "      <td>9.5</td>\n",
+       "      <td>608.1</td>\n",
+       "      <td>0.50</td>\n",
+       "      <td>0.07</td>\n",
+       "      <td>20.5</td>\n",
+       "      <td>-52.5</td>\n",
+       "      <td>521.6</td>\n",
+       "      <td>-1.08</td>\n",
+       "      <td>0.58</td>\n",
+       "      <td>-0.02</td>\n",
+       "      <td>-3.2</td>\n",
+       "      <td>103.6</td>\n",
+       "      <td>-0.95</td>\n",
+       "      <td>0.23</td>\n",
+       "      <td>-2.87</td>\n",
+       "      <td>-25.9</td>\n",
+       "      <td>-52.2</td>\n",
+       "      <td>-0.21</td>\n",
+       "      <td>0.87</td>\n",
+       "      <td>-1.81</td>\n",
+       "      <td>10.4</td>\n",
+       "      <td>62.0</td>\n",
+       "      <td>-0.28</td>\n",
+       "      <td>-0.04</td>\n",
+       "      <td>1.48</td>\n",
+       "      <td>-17.6</td>\n",
+       "      <td>-198.3</td>\n",
+       "      <td>3.43</td>\n",
+       "      <td>2.84</td>\n",
+       "      <td>5.87</td>\n",
+       "      <td>-16.9</td>\n",
+       "      <td>72.6</td>\n",
+       "      <td>-0.31</td>\n",
+       "      <td>2.79</td>\n",
+       "      <td>2.71</td>\n",
+       "      <td>-33.5</td>\n",
+       "      <td>-11.6</td>\n",
+       "      <td>-1.11</td>\n",
+       "      <td>4.01</td>\n",
+       "      <td>5.0</td>\n",
+       "      <td>-57.0</td>\n",
+       "      <td>666.3</td>\n",
+       "      <td>1.13</td>\n",
+       "      <td>4.38</td>\n",
+       "      <td>5.0</td>\n",
+       "      <td>-64.0</td>\n",
+       "      <td>39.3</td>\n",
+       "      <td>1.07</td>\n",
+       "      <td>-0.16</td>\n",
+       "      <td>32.5</td>\n",
+       "      <td>100.0</td>\n",
+       "      <td>1893.7</td>\n",
+       "      <td>-2.80</td>\n",
+       "      <td>-0.22</td>\n",
+       "      <td>2.5</td>\n",
+       "      <td>-28.5</td>\n",
+       "      <td>45.0</td>\n",
+       "      <td>0.58</td>\n",
+       "      <td>0.41</td>\n",
+       "      <td>-19.0</td>\n",
+       "      <td>-6.0</td>\n",
+       "      <td>762.9</td>\n",
+       "      <td>0.29</td>\n",
+       "      <td>0.82</td>\n",
+       "      <td>-3.0</td>\n",
+       "      <td>-35.0</td>\n",
+       "      <td>140.3</td>\n",
+       "      <td>1.16</td>\n",
+       "      <td>0.39</td>\n",
+       "      <td>0.73</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>2</th>\n",
+       "      <td>279</td>\n",
+       "      <td>261534</td>\n",
+       "      <td>0</td>\n",
+       "      <td>77.0</td>\n",
+       "      <td>27.27</td>\n",
+       "      <td>-0.91</td>\n",
+       "      <td>6.0</td>\n",
+       "      <td>58.5</td>\n",
+       "      <td>1623.6</td>\n",
+       "      <td>-1.40</td>\n",
+       "      <td>0.02</td>\n",
+       "      <td>-6.5</td>\n",
+       "      <td>-48.0</td>\n",
+       "      <td>621.0</td>\n",
+       "      <td>-1.20</td>\n",
+       "      <td>0.14</td>\n",
+       "      <td>-0.20</td>\n",
+       "      <td>73.6</td>\n",
+       "      <td>609.1</td>\n",
+       "      <td>-0.44</td>\n",
+       "      <td>-0.58</td>\n",
+       "      <td>-0.04</td>\n",
+       "      <td>-23.0</td>\n",
+       "      <td>-27.4</td>\n",
+       "      <td>-0.72</td>\n",
+       "      <td>-1.04</td>\n",
+       "      <td>-1.09</td>\n",
+       "      <td>91.1</td>\n",
+       "      <td>635.6</td>\n",
+       "      <td>-0.88</td>\n",
+       "      <td>0.24</td>\n",
+       "      <td>0.59</td>\n",
+       "      <td>-18.7</td>\n",
+       "      <td>-7.2</td>\n",
+       "      <td>-0.60</td>\n",
+       "      <td>-2.82</td>\n",
+       "      <td>-0.71</td>\n",
+       "      <td>52.4</td>\n",
+       "      <td>504.1</td>\n",
+       "      <td>0.89</td>\n",
+       "      <td>-0.67</td>\n",
+       "      <td>-9.30</td>\n",
+       "      <td>-20.8</td>\n",
+       "      <td>-25.7</td>\n",
+       "      <td>-0.77</td>\n",
+       "      <td>-0.85</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>-20.0</td>\n",
+       "      <td>2259.0</td>\n",
+       "      <td>-0.94</td>\n",
+       "      <td>1.15</td>\n",
+       "      <td>-4.0</td>\n",
+       "      <td>-44.0</td>\n",
+       "      <td>-22.7</td>\n",
+       "      <td>0.94</td>\n",
+       "      <td>-0.98</td>\n",
+       "      <td>-19.0</td>\n",
+       "      <td>105.0</td>\n",
+       "      <td>1267.9</td>\n",
+       "      <td>1.03</td>\n",
+       "      <td>1.27</td>\n",
+       "      <td>11.0</td>\n",
+       "      <td>-39.5</td>\n",
+       "      <td>82.3</td>\n",
+       "      <td>0.47</td>\n",
+       "      <td>-0.19</td>\n",
+       "      <td>-10.0</td>\n",
+       "      <td>7.0</td>\n",
+       "      <td>1491.8</td>\n",
+       "      <td>0.32</td>\n",
+       "      <td>-1.29</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>-34.0</td>\n",
+       "      <td>658.2</td>\n",
+       "      <td>-0.76</td>\n",
+       "      <td>0.26</td>\n",
+       "      <td>0.24</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>3</th>\n",
+       "      <td>279</td>\n",
+       "      <td>261535</td>\n",
+       "      <td>0</td>\n",
+       "      <td>41.0</td>\n",
+       "      <td>27.91</td>\n",
+       "      <td>-0.35</td>\n",
+       "      <td>3.0</td>\n",
+       "      <td>46.0</td>\n",
+       "      <td>1921.6</td>\n",
+       "      <td>-1.36</td>\n",
+       "      <td>-0.47</td>\n",
+       "      <td>-32.0</td>\n",
+       "      <td>-51.5</td>\n",
+       "      <td>560.9</td>\n",
+       "      <td>-0.29</td>\n",
+       "      <td>-0.10</td>\n",
+       "      <td>-1.11</td>\n",
+       "      <td>124.3</td>\n",
+       "      <td>791.6</td>\n",
+       "      <td>0.00</td>\n",
+       "      <td>0.39</td>\n",
+       "      <td>-1.85</td>\n",
+       "      <td>-21.7</td>\n",
+       "      <td>-44.9</td>\n",
+       "      <td>-0.21</td>\n",
+       "      <td>0.02</td>\n",
+       "      <td>0.89</td>\n",
+       "      <td>133.9</td>\n",
+       "      <td>797.8</td>\n",
+       "      <td>-0.08</td>\n",
+       "      <td>1.06</td>\n",
+       "      <td>-0.26</td>\n",
+       "      <td>-16.4</td>\n",
+       "      <td>-74.1</td>\n",
+       "      <td>0.97</td>\n",
+       "      <td>-0.80</td>\n",
+       "      <td>-0.41</td>\n",
+       "      <td>66.9</td>\n",
+       "      <td>955.3</td>\n",
+       "      <td>-1.90</td>\n",
+       "      <td>1.28</td>\n",
+       "      <td>-6.65</td>\n",
+       "      <td>-28.1</td>\n",
+       "      <td>47.5</td>\n",
+       "      <td>-1.91</td>\n",
+       "      <td>1.42</td>\n",
+       "      <td>1.0</td>\n",
+       "      <td>-30.0</td>\n",
+       "      <td>1846.7</td>\n",
+       "      <td>0.76</td>\n",
+       "      <td>1.10</td>\n",
+       "      <td>-4.0</td>\n",
+       "      <td>-52.0</td>\n",
+       "      <td>-53.9</td>\n",
+       "      <td>1.71</td>\n",
+       "      <td>-0.22</td>\n",
+       "      <td>-12.0</td>\n",
+       "      <td>97.5</td>\n",
+       "      <td>1969.8</td>\n",
+       "      <td>-1.70</td>\n",
+       "      <td>0.16</td>\n",
+       "      <td>-1.0</td>\n",
+       "      <td>-32.5</td>\n",
+       "      <td>255.9</td>\n",
+       "      <td>-0.46</td>\n",
+       "      <td>1.57</td>\n",
+       "      <td>10.0</td>\n",
+       "      <td>6.0</td>\n",
+       "      <td>2047.7</td>\n",
+       "      <td>-0.98</td>\n",
+       "      <td>1.53</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>-49.0</td>\n",
+       "      <td>554.2</td>\n",
+       "      <td>-0.83</td>\n",
+       "      <td>0.39</td>\n",
+       "      <td>0.73</td>\n",
+       "    </tr>\n",
+       "    <tr>\n",
+       "      <th>4</th>\n",
+       "      <td>279</td>\n",
+       "      <td>261536</td>\n",
+       "      <td>0</td>\n",
+       "      <td>50.0</td>\n",
+       "      <td>28.00</td>\n",
+       "      <td>-1.32</td>\n",
+       "      <td>-9.0</td>\n",
+       "      <td>12.0</td>\n",
+       "      <td>464.8</td>\n",
+       "      <td>0.88</td>\n",
+       "      <td>0.19</td>\n",
+       "      <td>8.0</td>\n",
+       "      <td>-51.5</td>\n",
+       "      <td>98.1</td>\n",
+       "      <td>1.09</td>\n",
+       "      <td>-0.33</td>\n",
+       "      <td>-2.16</td>\n",
+       "      <td>-3.9</td>\n",
+       "      <td>102.7</td>\n",
+       "      <td>0.39</td>\n",
+       "      <td>-1.22</td>\n",
+       "      <td>-3.39</td>\n",
+       "      <td>-15.2</td>\n",
+       "      <td>-42.2</td>\n",
+       "      <td>-1.18</td>\n",
+       "      <td>-1.11</td>\n",
+       "      <td>-3.55</td>\n",
+       "      <td>8.9</td>\n",
+       "      <td>141.3</td>\n",
+       "      <td>-0.16</td>\n",
+       "      <td>-0.43</td>\n",
+       "      <td>-4.15</td>\n",
+       "      <td>-12.9</td>\n",
+       "      <td>-13.4</td>\n",
+       "      <td>-1.32</td>\n",
+       "      <td>-0.98</td>\n",
+       "      <td>-3.69</td>\n",
+       "      <td>8.8</td>\n",
+       "      <td>136.1</td>\n",
+       "      <td>-0.30</td>\n",
+       "      <td>4.13</td>\n",
+       "      <td>1.89</td>\n",
+       "      <td>-13.0</td>\n",
+       "      <td>-18.7</td>\n",
+       "      <td>-1.37</td>\n",
+       "      <td>-0.93</td>\n",
+       "      <td>0.0</td>\n",
+       "      <td>-1.0</td>\n",
+       "      <td>810.1</td>\n",
+       "      <td>-2.29</td>\n",
+       "      <td>6.72</td>\n",
+       "      <td>1.0</td>\n",
+       "      <td>-23.0</td>\n",
+       "      <td>-29.7</td>\n",
+       "      <td>0.58</td>\n",
+       "      <td>-1.10</td>\n",
+       "      <td>-18.5</td>\n",
+       "      <td>33.5</td>\n",
+       "      <td>206.8</td>\n",
+       "      <td>1.84</td>\n",
+       "      <td>-0.13</td>\n",
+       "      <td>4.0</td>\n",
+       "      <td>-29.0</td>\n",
+       "      <td>30.1</td>\n",
+       "      <td>0.80</td>\n",
+       "      <td>-0.24</td>\n",
+       "      <td>5.0</td>\n",
+       "      <td>-14.0</td>\n",
+       "      <td>479.5</td>\n",
+       "      <td>0.68</td>\n",
+       "      <td>-0.59</td>\n",
+       "      <td>2.0</td>\n",
+       "      <td>-36.0</td>\n",
+       "      <td>-6.9</td>\n",
+       "      <td>2.02</td>\n",
+       "      <td>0.14</td>\n",
+       "      <td>-0.23</td>\n",
+       "    </tr>\n",
+       "  </tbody>\n",
+       "</table>\n",
+       "</div>"
+      ],
+      "text/plain": [
+       "   EXAMPLE_ID  BLOCK_ID  target     4      5     6     7     8       9    10  \\\n",
+       "0         279    261532       0  52.0  32.69  0.30   2.5  20.0  1256.8 -0.89   \n",
+       "1         279    261533       0  58.0  33.33  0.00  16.5   9.5   608.1  0.50   \n",
+       "2         279    261534       0  77.0  27.27 -0.91   6.0  58.5  1623.6 -1.40   \n",
+       "3         279    261535       0  41.0  27.91 -0.35   3.0  46.0  1921.6 -1.36   \n",
+       "4         279    261536       0  50.0  28.00 -1.32  -9.0  12.0   464.8  0.88   \n",
+       "\n",
+       "     11    12    13     14    15    16    17     18     19    20    21    22  \\\n",
+       "0  0.33  11.0 -55.0  267.2  0.52  0.05 -2.36   49.6  252.0  0.43  1.16 -2.06   \n",
+       "1  0.07  20.5 -52.5  521.6 -1.08  0.58 -0.02   -3.2  103.6 -0.95  0.23 -2.87   \n",
+       "2  0.02  -6.5 -48.0  621.0 -1.20  0.14 -0.20   73.6  609.1 -0.44 -0.58 -0.04   \n",
+       "3 -0.47 -32.0 -51.5  560.9 -0.29 -0.10 -1.11  124.3  791.6  0.00  0.39 -1.85   \n",
+       "4  0.19   8.0 -51.5   98.1  1.09 -0.33 -2.16   -3.9  102.7  0.39 -1.22 -3.39   \n",
+       "\n",
+       "     23     24    25    26    27     28     29    30    31    32    33     34  \\\n",
+       "0 -33.0 -123.2  1.60 -0.49 -6.06   65.0  296.1 -0.28 -0.26 -3.83 -22.6 -170.0   \n",
+       "1 -25.9  -52.2 -0.21  0.87 -1.81   10.4   62.0 -0.28 -0.04  1.48 -17.6 -198.3   \n",
+       "2 -23.0  -27.4 -0.72 -1.04 -1.09   91.1  635.6 -0.88  0.24  0.59 -18.7   -7.2   \n",
+       "3 -21.7  -44.9 -0.21  0.02  0.89  133.9  797.8 -0.08  1.06 -0.26 -16.4  -74.1   \n",
+       "4 -15.2  -42.2 -1.18 -1.11 -3.55    8.9  141.3 -0.16 -0.43 -4.15 -12.9  -13.4   \n",
+       "\n",
+       "     35    36    37    38     39    40    41    42    43     44    45    46  \\\n",
+       "0  3.06 -1.05 -3.29  22.9  286.3  0.12  2.58  4.08 -33.0 -178.9  1.88  0.53   \n",
+       "1  3.43  2.84  5.87 -16.9   72.6 -0.31  2.79  2.71 -33.5  -11.6 -1.11  4.01   \n",
+       "2 -0.60 -2.82 -0.71  52.4  504.1  0.89 -0.67 -9.30 -20.8  -25.7 -0.77 -0.85   \n",
+       "3  0.97 -0.80 -0.41  66.9  955.3 -1.90  1.28 -6.65 -28.1   47.5 -1.91  1.42   \n",
+       "4 -1.32 -0.98 -3.69   8.8  136.1 -0.30  4.13  1.89 -13.0  -18.7 -1.37 -0.93   \n",
+       "\n",
+       "    47    48      49    50    51   52    53     54    55    56    57     58  \\\n",
+       "0 -7.0 -44.0  1987.0 -5.41  0.95 -4.0 -57.0  722.9 -3.26 -0.55  -7.5  125.5   \n",
+       "1  5.0 -57.0   666.3  1.13  4.38  5.0 -64.0   39.3  1.07 -0.16  32.5  100.0   \n",
+       "2  0.0 -20.0  2259.0 -0.94  1.15 -4.0 -44.0  -22.7  0.94 -0.98 -19.0  105.0   \n",
+       "3  1.0 -30.0  1846.7  0.76  1.10 -4.0 -52.0  -53.9  1.71 -0.22 -12.0   97.5   \n",
+       "4  0.0  -1.0   810.1 -2.29  6.72  1.0 -23.0  -29.7  0.58 -1.10 -18.5   33.5   \n",
+       "\n",
+       "       59    60    61    62    63     64    65    66    67    68      69  \\\n",
+       "0  1547.2 -0.36  1.12   9.0 -37.0   72.5  0.47  0.74 -11.0  -8.0  1595.1   \n",
+       "1  1893.7 -2.80 -0.22   2.5 -28.5   45.0  0.58  0.41 -19.0  -6.0   762.9   \n",
+       "2  1267.9  1.03  1.27  11.0 -39.5   82.3  0.47 -0.19 -10.0   7.0  1491.8   \n",
+       "3  1969.8 -1.70  0.16  -1.0 -32.5  255.9 -0.46  1.57  10.0   6.0  2047.7   \n",
+       "4   206.8  1.84 -0.13   4.0 -29.0   30.1  0.80 -0.24   5.0 -14.0   479.5   \n",
+       "\n",
+       "     70    71   72    73     74    75    76    77  \n",
+       "0 -1.64  2.83 -2.0 -50.0  445.2 -0.35  0.26  0.76  \n",
+       "1  0.29  0.82 -3.0 -35.0  140.3  1.16  0.39  0.73  \n",
+       "2  0.32 -1.29  0.0 -34.0  658.2 -0.76  0.26  0.24  \n",
+       "3 -0.98  1.53  0.0 -49.0  554.2 -0.83  0.39  0.73  \n",
+       "4  0.68 -0.59  2.0 -36.0   -6.9  2.02  0.14 -0.23  "
+      ]
+     },
+     "execution_count": 2,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "header_list = ['EXAMPLE_ID', 'BLOCK_ID', 'target'] + [str(i) for i in range(4,78)]\n",
+    "df = pd.read_csv('data/kddcup04/bio_train.dat', sep='\\t', names=header_list)\n",
+    "df.head()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "0    144455\n",
+       "1      1296\n",
+       "Name: target, dtype: int64"
+      ]
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "# imbalance of the classes\n",
+    "df['target'].value_counts()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# drop columns we won't need in this example\n",
+    "df.drop(columns=['EXAMPLE_ID', 'BLOCK_ID'], inplace=True)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "df_train, df_valid = train_test_split(df, test_size=0.2, stratify=df['target'], random_state=1)\n",
+    "df_valid, df_test = train_test_split(df_valid, test_size=0.5, stratify=df_valid['target'], random_state=1)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Preparing the data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "continuous_cols = df.drop(columns=['target']).columns.values.tolist()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# deeptabular\n",
+    "tab_preprocessor = TabPreprocessor(continuous_cols=continuous_cols, scale=True)\n",
+    "X_tab_train = tab_preprocessor.fit_transform(df_train)\n",
+    "X_tab_valid = tab_preprocessor.transform(df_valid)\n",
+    "X_tab_test = tab_preprocessor.transform(df_test)\n",
+    "\n",
+    "# target\n",
+    "y_train = df_train[\"target\"].values\n",
+    "y_valid = df_valid[\"target\"].values\n",
+    "y_test = df_test[\"target\"].values"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Define the model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "input_layer = len(tab_preprocessor.continuous_cols)\n",
+    "output_layer = 1\n",
+    "hidden_layers = np.linspace(\n",
+    "    input_layer * 2, output_layer, 5, endpoint=False, dtype=int\n",
+    ").tolist()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "WideDeep(\n",
+       "  (deeptabular): Sequential(\n",
+       "    (0): TabMlp(\n",
+       "      (cat_embed_and_cont): CatEmbeddingsAndCont(\n",
+       "        (cont_norm): BatchNorm1d(74, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "      (tab_mlp): MLP(\n",
+       "        (mlp): Sequential(\n",
+       "          (dense_layer_0): Sequential(\n",
+       "            (0): Dropout(p=0.1, inplace=False)\n",
+       "            (1): Linear(in_features=74, out_features=148, bias=True)\n",
+       "            (2): ReLU(inplace=True)\n",
+       "          )\n",
+       "          (dense_layer_1): Sequential(\n",
+       "            (0): Dropout(p=0.1, inplace=False)\n",
+       "            (1): Linear(in_features=148, out_features=118, bias=True)\n",
+       "            (2): ReLU(inplace=True)\n",
+       "          )\n",
+       "          (dense_layer_2): Sequential(\n",
+       "            (0): Dropout(p=0.1, inplace=False)\n",
+       "            (1): Linear(in_features=118, out_features=89, bias=True)\n",
+       "            (2): ReLU(inplace=True)\n",
+       "          )\n",
+       "          (dense_layer_3): Sequential(\n",
+       "            (0): Dropout(p=0.1, inplace=False)\n",
+       "            (1): Linear(in_features=89, out_features=59, bias=True)\n",
+       "            (2): ReLU(inplace=True)\n",
+       "          )\n",
+       "          (dense_layer_4): Sequential(\n",
+       "            (0): Dropout(p=0.1, inplace=False)\n",
+       "            (1): Linear(in_features=59, out_features=30, bias=True)\n",
+       "            (2): ReLU(inplace=True)\n",
+       "          )\n",
+       "        )\n",
+       "      )\n",
+       "    )\n",
+       "    (1): Linear(in_features=30, out_features=1, bias=True)\n",
+       "  )\n",
+       ")"
+      ]
+     },
+     "execution_count": 9,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "deeptabular = TabMlp(\n",
+    "    mlp_hidden_dims=hidden_layers,\n",
+    "    column_idx=tab_preprocessor.column_idx,\n",
+    "    continuous_cols=tab_preprocessor.continuous_cols,\n",
+    ")\n",
+    "model = WideDeep(deeptabular=deeptabular)\n",
+    "model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Metrics from torchmetrics\n",
+    "accuracy = Accuracy_torchmetrics(average=None, num_classes=2)\n",
+    "precision = Precision_torchmetrics(average=\"micro\", num_classes=2)\n",
+    "f1 = F1_torchmetrics(average=None, num_classes=2)\n",
+    "recall = Recall_torchmetrics(average=None, num_classes=2)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# # Metrics from pytorch-widedeep\n",
+    "# accuracy = Accuracy(top_k=2)\n",
+    "# precision = Precision(average=False)\n",
+    "# recall = Recall(average=True)\n",
+    "# f1 = F1Score(average=False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 33,
+   "metadata": {
+    "scrolled": false
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/html": [
+       "== Status ==<br>Memory usage on this node: 10.5/16.0 GiB<br>Using FIFO scheduling algorithm.<br>Resources requested: 0/8 CPUs, 0/0 GPUs, 0.0/4.32 GiB heap, 0.0/2.16 GiB objects<br>Result logdir: /Users/javier/ray_results/training_function_2021-10-18_15-59-06<br>Number of trials: 2/2 (2 TERMINATED)<br><table>\n",
+       "<thead>\n",
+       "<tr><th>Trial name                   </th><th>status    </th><th>loc  </th><th style=\"text-align: right;\">  batch_size</th><th style=\"text-align: right;\">  iter</th><th style=\"text-align: right;\">  total time (s)</th></tr>\n",
+       "</thead>\n",
+       "<tbody>\n",
+       "<tr><td>training_function_8f035_00000</td><td>TERMINATED</td><td>     </td><td style=\"text-align: right;\">        1000</td><td style=\"text-align: right;\">     5</td><td style=\"text-align: right;\">         18.2589</td></tr>\n",
+       "<tr><td>training_function_8f035_00001</td><td>TERMINATED</td><td>     </td><td style=\"text-align: right;\">        5000</td><td style=\"text-align: right;\">     5</td><td style=\"text-align: right;\">         18.0369</td></tr>\n",
+       "</tbody>\n",
+       "</table><br><br>"
+      ],
+      "text/plain": [
+       "<IPython.core.display.HTML object>"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "2021-10-18 15:59:28,373\tINFO tune.py:617 -- Total run time: 22.07 seconds (21.91 seconds for the tuning loop).\n"
+     ]
+    }
+   ],
+   "source": [
+    "config = {\n",
+    "    \"batch_size\": tune.grid_search([1000, 5000]),\n",
+    "}\n",
+    "\n",
+    "# Optimizers\n",
+    "deep_opt = SGD(model.deeptabular.parameters(), lr=0.1)\n",
+    "# LR Scheduler\n",
+    "deep_sch = lr_scheduler.StepLR(deep_opt, step_size=3)\n",
+    "\n",
+    "early_stopping = EarlyStopping()\n",
+    "\n",
+    "\n",
+    "def training_function(config, X_train, X_val):\n",
+    "    # Hyperparameters\n",
+    "    batch_size = config[\"batch_size\"]\n",
+    "    trainer = Trainer(\n",
+    "        model,\n",
+    "        objective=\"binary_focal_loss\",\n",
+    "        callbacks=[RayTuneReporter, LRHistory(n_epochs=10), early_stopping],\n",
+    "        lr_schedulers={\"deeptabular\": deep_sch},\n",
+    "        initializers={\"deeptabular\": XavierNormal},\n",
+    "        optimizers={\"deeptabular\": deep_opt},\n",
+    "        metrics=[accuracy, precision, recall, f1],\n",
+    "        verbose=0,\n",
+    "    )\n",
+    "\n",
+    "    trainer.fit(X_train=X_train, X_val=X_val, n_epochs=5, batch_size=batch_size)\n",
+    "\n",
+    "\n",
+    "X_train = {\"X_tab\": X_tab_train, \"target\": y_train}\n",
+    "X_val = {\"X_tab\": X_tab_valid, \"target\": y_valid}\n",
+    "analysis = tune.run(\n",
+    "    tune.with_parameters(training_function, X_train=X_train, X_val=X_val),\n",
+    "    resources_per_trial={\"cpu\": 1, \"gpu\": 0},\n",
+    "    progress_reporter=JupyterNotebookReporter(overwrite=True),\n",
+    "    config=config,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 34,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "{'8f035_00000': {'_metric': {'train_loss': 0.007156053689332345,\n",
+       "   'train_Accuracy_0': 1.0,\n",
+       "   'train_Accuracy_1': 0.008678881451487541,\n",
+       "   'train_Precision': 0.9911835193634033,\n",
+       "   'train_Recall_0': 1.0,\n",
+       "   'train_Recall_1': 0.008678881451487541,\n",
+       "   'train_F1_0': 0.9955719113349915,\n",
+       "   'train_F1_1': 0.017208412289619446,\n",
+       "   'val_loss': 0.006261252580831448,\n",
+       "   'val_Accuracy_0': 1.0,\n",
+       "   'val_Accuracy_1': 0.023255813866853714,\n",
+       "   'val_Precision': 0.9913550615310669,\n",
+       "   'val_Recall_0': 1.0,\n",
+       "   'val_Recall_1': 0.023255813866853714,\n",
+       "   'val_F1_0': 0.9956578612327576,\n",
+       "   'val_F1_1': 0.045454543083906174,\n",
+       "   'lr_deeptabular_0': 0.010000000000000002},\n",
+       "  'time_this_iter_s': 3.5364139080047607,\n",
+       "  'done': True,\n",
+       "  'timesteps_total': None,\n",
+       "  'episodes_total': None,\n",
+       "  'training_iteration': 5,\n",
+       "  'experiment_id': 'f62bb9c9c32a45b9af85a6a3b0e30d94',\n",
+       "  'date': '2021-10-18_15-59-28',\n",
+       "  'timestamp': 1634565568,\n",
+       "  'time_total_s': 18.25888180732727,\n",
+       "  'pid': 15457,\n",
+       "  'hostname': 'infinito.bbrouter',\n",
+       "  'node_ip': '192.168.18.34',\n",
+       "  'config': {'batch_size': 1000},\n",
+       "  'time_since_restore': 18.25888180732727,\n",
+       "  'timesteps_since_restore': 0,\n",
+       "  'iterations_since_restore': 5,\n",
+       "  'trial_id': '8f035_00000',\n",
+       "  'experiment_tag': '0_batch_size=1000'},\n",
+       " '8f035_00001': {'_metric': {'train_loss': 0.019367387828727562,\n",
+       "   'train_Accuracy_0': 0.9999827146530151,\n",
+       "   'train_Accuracy_1': 0.01157184224575758,\n",
+       "   'train_Precision': 0.991192102432251,\n",
+       "   'train_Recall_0': 0.9999827146530151,\n",
+       "   'train_Recall_1': 0.01157184224575758,\n",
+       "   'train_F1_0': 0.9955761432647705,\n",
+       "   'train_F1_1': 0.022835396230220795,\n",
+       "   'val_loss': 0.01834123209118843,\n",
+       "   'val_Accuracy_0': 1.0,\n",
+       "   'val_Accuracy_1': 0.0,\n",
+       "   'val_Precision': 0.9911492466926575,\n",
+       "   'val_Recall_0': 1.0,\n",
+       "   'val_Recall_1': 0.0,\n",
+       "   'val_F1_0': 0.9955549836158752,\n",
+       "   'val_F1_1': 0.0,\n",
+       "   'lr_deeptabular_0': 0.010000000000000002},\n",
+       "  'time_this_iter_s': 3.42478084564209,\n",
+       "  'done': True,\n",
+       "  'timesteps_total': None,\n",
+       "  'episodes_total': None,\n",
+       "  'training_iteration': 5,\n",
+       "  'experiment_id': 'd7379f1debc14e41bc971bf4c27b6793',\n",
+       "  'date': '2021-10-18_15-59-27',\n",
+       "  'timestamp': 1634565567,\n",
+       "  'time_total_s': 18.036858797073364,\n",
+       "  'pid': 15456,\n",
+       "  'hostname': 'infinito.bbrouter',\n",
+       "  'node_ip': '192.168.18.34',\n",
+       "  'config': {'batch_size': 5000},\n",
+       "  'time_since_restore': 18.036858797073364,\n",
+       "  'timesteps_since_restore': 0,\n",
+       "  'iterations_since_restore': 5,\n",
+       "  'trial_id': '8f035_00001',\n",
+       "  'experiment_tag': '1_batch_size=5000'}}"
+      ]
+     },
+     "execution_count": 34,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "analysis.results"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# %load_ext tensorboard"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [],
+   "source": [
+    "# %tensorboard --logdir ~/ray_results"
+   ]
+  }
+ ],
+ "metadata": {
+  "interpreter": {
+   "hash": "3b99005fd577fa40f3cce433b2b92303885900e634b2b5344c07c59d06c8792d"
+  },
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.7.7"
+  },
+  "toc": {
+   "base_numbering": 1,
+   "nav_menu": {},
+   "number_sections": true,
+   "sideBar": true,
+   "skip_h1_title": false,
+   "title_cell": "Table of Contents",
+   "title_sidebar": "Contents",
+   "toc_cell": false,
+   "toc_position": {},
+   "toc_section_display": true,
+   "toc_window_display": false
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}
diff --git a/examples/adult_census.py b/examples/adult_census.py
index 6656e489..76597d85 100644
--- a/examples/adult_census.py
+++ b/examples/adult_census.py
@@ -42,12 +42,19 @@
         "gender",
     ]
     crossed_cols = [("education", "occupation"), ("native_country", "occupation")]
+    # cat_embed_cols = [
+    #     ("education", 10),
+    #     ("relationship", 8),
+    #     ("workclass", 10),
+    #     ("occupation", 10),
+    #     ("native_country", 10),
+    # ]
     cat_embed_cols = [
-        ("education", 10),
-        ("relationship", 8),
-        ("workclass", 10),
-        ("occupation", 10),
-        ("native_country", 10),
+        "education",
+        "relationship",
+        "workclass",
+        "occupation",
+        "native_country",
     ]
     continuous_cols = ["age", "hours_per_week"]
     target = "income_label"
diff --git a/pytorch_widedeep/callbacks.py b/pytorch_widedeep/callbacks.py
index f39b6d57..b2a338f8 100644
--- a/pytorch_widedeep/callbacks.py
+++ b/pytorch_widedeep/callbacks.py
@@ -9,6 +9,7 @@
 
 import numpy as np
 import torch
+from ray import tune
 from torch.optim.lr_scheduler import ReduceLROnPlateau
 
 from pytorch_widedeep.metrics import MultipleMetrics
@@ -159,7 +160,11 @@ def on_epoch_end(
         for k, v in logs.items():
             if isinstance(v, np.ndarray):
                 v = v.tolist()
-            self.trainer.history.setdefault(k, []).append(v)
+            if isinstance(v, list) and len(v) > 1:
+                for i in range(len(v)):
+                    self.trainer.history.setdefault(k + "_" + str(i), []).append(v[i])
+            else:
+                self.trainer.history.setdefault(k, []).append(v)
 
 
 class LRShedulerCallback(Callback):
@@ -661,3 +666,27 @@ def __getstate__(self):
 
     def __setstate__(self, state):
         self.__dict__ = state
+
+
+class RayTuneReporter(Callback):
+    r"""Callback that allows reporting history and lr_history values to RayTune
+    during Hyperparameter tuning
+
+    Callbacks are passed as input parameters to the :obj:`Trainer` class. See
+    :class:`pytorch_widedeep.trainer.Trainer`
+
+    Examples
+    --------
+    see /examples/12_HyperParameter_tuning_w_RayTune.ipynb
+    """
+
+    def on_epoch_end(
+        self, epoch: int, logs: Optional[Dict] = None, metric: Optional[float] = None
+    ):
+        report_dict = {}
+        for k, v in self.trainer.history.items():
+            report_dict.update({k: v[-1]})
+        if hasattr(self.trainer, "lr_history"):
+            for k, v in self.trainer.lr_history.items():
+                report_dict.update({k: v[-1]})
+        tune.report(report_dict)
diff --git a/pytorch_widedeep/tab2vec.py b/pytorch_widedeep/tab2vec.py
index bd46305a..c390143a 100644
--- a/pytorch_widedeep/tab2vec.py
+++ b/pytorch_widedeep/tab2vec.py
@@ -32,6 +32,9 @@ class Tab2Vec:
     tab_preprocessor: ``TabPreprocessor``
         ``TabPreprocessor`` object. Must be fitted.
         See :obj:`pytorch-widedeep.preprocessing.tab_preprocessor.TabPreprocessor`
+    return_dataframe: bool
+        Boolean indicating of the returned object(s) will be array(s) or
+        pandas dataframe(s)
 
     Attributes
     ----------
@@ -81,7 +84,11 @@ class Tab2Vec:
     """
 
     def __init__(
-        self, model: WideDeep, tab_preprocessor: TabPreprocessor, verbose: bool = False
+        self,
+        model: WideDeep,
+        tab_preprocessor: TabPreprocessor,
+        return_dataframe: bool = False,
+        verbose: bool = False,
     ):
         super(Tab2Vec, self).__init__()
 
@@ -100,6 +107,7 @@ def __init__(
                 "The 'tab_preprocessor' must be fitted before is passed to 'Tab2Vec'"
             )
 
+        self.return_dataframe = return_dataframe
         self.tab_preprocessor = tab_preprocessor
 
         transformer_family = [
@@ -126,8 +134,15 @@ def fit(self, df: pd.DataFrame, target_col: Optional[str] = None) -> "Tab2Vec":
         return self
 
     def transform(
-        self, df: pd.DataFrame, target_col: Optional[str] = None
-    ) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
+        self,
+        df: pd.DataFrame,
+        target_col: Optional[str] = None,
+    ) -> Union[
+        np.ndarray,
+        Tuple[np.ndarray, np.ndarray],
+        pd.DataFrame,
+        Tuple[pd.DataFrame, pd.Series],
+    ]:
         r"""
         Parameters
         ----------
@@ -155,13 +170,60 @@ def transform(
                 x_cont = einops.rearrange(x_cont, "s c e -> s (c e)")
 
         X_vec = torch.cat([x_cat, x_cont], 1).cpu().numpy()
-        if target_col:
-            return X_vec, df[target_col].values
+
+        if self.return_dataframe:
+            new_colnames = self._new_colnames()
+            if target_col:
+                return pd.DataFrame(data=X_vec, columns=new_colnames), df[[target_col]]
+            else:
+                return pd.DataFrame(data=X_vec, columns=new_colnames)
         else:
-            return X_vec
+            if target_col:
+                return X_vec, df[target_col].values
+            else:
+                return X_vec
 
     def fit_transform(
         self, df: pd.DataFrame, target_col: Optional[str] = None
     ) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
         r"""Combines ``fit`` and ``transform``"""
         return self.fit(df, target_col).transform(df, target_col)
+
+    def _new_colnames(self) -> List[str]:
+
+        if self.is_transformer:
+            new_colnames = []
+            embedding_dim = self.vectorizer.cat_embed.embed.embedding_dim
+            are_cont_embed = hasattr(self.vectorizer, "cont_embed")
+            cat_cols = [
+                ei[0]
+                for ei in self.tab_preprocessor.embeddings_input
+                if ei[0] != "cls_token"
+            ]
+            all_embed_cols = (
+                cat_cols + self.tab_preprocessor.continuous_cols
+                if are_cont_embed
+                else cat_cols
+            )
+            for colname in all_embed_cols:
+                new_colnames.extend(
+                    [
+                        "_".join([colname, "embed", str(i + 1)])
+                        for i in range(embedding_dim)
+                    ]
+                )
+            if not are_cont_embed:
+                new_colnames += self.tab_preprocessor.continuous_cols
+        else:
+            new_colnames = []
+            for colname, _, embedding_dim in self.tab_preprocessor.embeddings_input:
+                new_colnames.extend(
+                    [
+                        "_".join([colname, "embed", str(i + 1)])
+                        for i in range(embedding_dim)
+                    ]
+                )
+            if self.tab_preprocessor.continuous_cols is not None:
+                new_colnames += self.tab_preprocessor.continuous_cols
+
+        return new_colnames
diff --git a/pytorch_widedeep/training/trainer.py b/pytorch_widedeep/training/trainer.py
index 4334ac9e..1225c9de 100644
--- a/pytorch_widedeep/training/trainer.py
+++ b/pytorch_widedeep/training/trainer.py
@@ -1,6 +1,7 @@
 import os
 import json
 import warnings
+from copy import deepcopy
 from pathlib import Path
 
 import numpy as np
@@ -19,6 +20,7 @@
     History,
     Callback,
     MetricCallback,
+    RayTuneReporter,
     CallbackContainer,
     LRShedulerCallback,
 )
diff --git a/pytorch_widedeep/version.py b/pytorch_widedeep/version.py
index 9fd0f8dd..9eb1ebec 100644
--- a/pytorch_widedeep/version.py
+++ b/pytorch_widedeep/version.py
@@ -1,1 +1,1 @@
-__version__ = "1.0.10"
+__version__ = "1.0.11"
diff --git a/setup.py b/setup.py
index ebe38678..2f726875 100644
--- a/setup.py
+++ b/setup.py
@@ -65,6 +65,7 @@
         "einops",
         "wrapt",
         "torchmetrics",
+        "ray[tune]",
     ],
     "extras_require": extras,
     "python_requires": ">=3.7.0",


Test patch:
diff --git a/tests/test_model_functioning/test_callbacks.py b/tests/test_model_functioning/test_callbacks.py
index e1793c1c..d36f98c9 100644
--- a/tests/test_model_functioning/test_callbacks.py
+++ b/tests/test_model_functioning/test_callbacks.py
@@ -8,6 +8,7 @@
 import numpy as np
 import torch
 import pytest
+from ray import tune
 from torch.optim.lr_scheduler import StepLR, CyclicLR
 
 from pytorch_widedeep.optim import RAdam
@@ -17,6 +18,7 @@
     LRHistory,
     EarlyStopping,
     ModelCheckpoint,
+    RayTuneReporter,
 )
 
 # Wide array
@@ -467,3 +469,42 @@ def test_early_stopping_get_state():
     shutil.rmtree("tests/test_model_functioning/early_stopping/")
 
     assert no_trainer and no_model
+
+
+###############################################################################
+# Test RayTuneReporter
+###############################################################################
+
+
+def test_ray_tune_reporter():
+
+    config = {
+        "batch_size": tune.grid_search([8, 16]),
+    }
+
+    def training_function(config):
+        batch_size = config["batch_size"]
+
+        trainer = Trainer(
+            model,
+            objective="binary",
+            callbacks=[RayTuneReporter],
+            verbose=0,
+        )
+
+        trainer.fit(
+            X_wide=X_wide,
+            X_tab=X_tab,
+            target=target,
+            n_epochs=1,
+            batch_size=batch_size,
+        )
+
+    analysis = tune.run(
+        tune.with_parameters(training_function),
+        config=config,
+        resources_per_trial={"cpu": 1, "gpu": 0},
+        verbose=0,
+    )
+
+    assert any(["train_loss" in el for el in analysis.results_df.keys()])
diff --git a/tests/test_tab2vec/test_t2v.py b/tests/test_tab2vec/test_t2v.py
index c2468fa8..aafa8720 100644
--- a/tests/test_tab2vec/test_t2v.py
+++ b/tests/test_tab2vec/test_t2v.py
@@ -27,7 +27,7 @@
 def create_df():
     cat_cols = [np.array(choices(c, k=5)) for c in [cat_col1_vals, cat_col2_vals]]
     cont_cols = [np.round(np.random.rand(5), 2) for _ in range(2)]
-    target = [np.random.choice(2, 5)]
+    target = [np.random.choice(2, 5, p=[0.8, 0.2])]
     return pd.DataFrame(
         np.vstack(cat_cols + cont_cols + target).transpose(), columns=colnames
     )
@@ -63,20 +63,27 @@ def create_df():
 
 
 @pytest.mark.parametrize(
-    "deeptabular",
-    [tabmlp, tabresnet, tabnet],
+    "deeptabular, return_dataframe",
+    [
+        (tabmlp, True),
+        (tabmlp, False),
+        (tabresnet, True),
+        (tabresnet, False),
+        (tabnet, True),
+        (tabnet, False),
+    ],
 )
-def test_non_transformer_models(deeptabular):
+def test_non_transformer_models(deeptabular, return_dataframe):
 
     model = WideDeep(deeptabular=deeptabular)
 
     # Let's assume the model is trained
-    t2v = Tab2Vec(model, tab_preprocessor)
-    X_vec, _ = t2v.fit_transform(df_t2v, target_col="target")
+    t2v = Tab2Vec(model, tab_preprocessor, return_dataframe=return_dataframe)
+    t2v_out, _ = t2v.fit_transform(df_t2v, target_col="target")
 
     embed_dim = sum([el[2] for el in tab_preprocessor.embeddings_input])
     cont_dim = len(tab_preprocessor.continuous_cols)
-    assert X_vec.shape[1] == embed_dim + cont_dim
+    assert t2v_out.shape[1] == embed_dim + cont_dim
 
 
 ###############################################################################
@@ -154,26 +161,45 @@ def test_tab_transformer_models(
 
 
 @pytest.mark.parametrize(
-    "model_name, with_cls_token, share_embeddings",
+    "model_name, with_cls_token, share_embeddings, return_dataframe",
     [
-        ("saint", False, True),
-        ("saint", True, True),
-        ("saint", False, False),
-        ("fttransformer", False, True),
-        ("fttransformer", True, True),
-        ("fttransformer", False, False),
-        ("tabfastformer", False, True),
-        ("tabfastformer", True, True),
-        ("tabfastformer", False, False),
+        ("saint", False, True, False),
+        ("saint", True, True, False),
+        ("saint", False, False, False),
+        ("saint", False, True, True),
+        ("saint", True, True, True),
+        ("saint", False, False, True),
+        ("fttransformer", False, True, False),
+        ("fttransformer", True, True, False),
+        ("fttransformer", False, False, False),
+        ("fttransformer", False, True, True),
+        ("fttransformer", True, True, True),
+        ("fttransformer", False, False, True),
+        ("tabfastformer", False, True, False),
+        ("tabfastformer", True, True, False),
+        ("tabfastformer", False, False, False),
+        ("tabfastformer", False, True, True),
+        ("tabfastformer", True, True, True),
+        ("tabfastformer", False, False, True),
         (
             "tabperceiver",
             False,
             True,
+            False,
         ),  # for the perceiver we do not need with_cls_token
-        ("tabperceiver", False, False),
+        ("tabperceiver", False, False, False),
+        (
+            "tabperceiver",
+            False,
+            True,
+            False,
+        ),
+        ("tabperceiver", False, False, True),
     ],
 )
-def test_transformer_family_models(model_name, with_cls_token, share_embeddings):
+def test_transformer_family_models(
+    model_name, with_cls_token, share_embeddings, return_dataframe
+):
 
     embed_cols = ["a", "b"]
     cont_cols = ["c", "d"]
@@ -197,9 +223,9 @@ def test_transformer_family_models(model_name, with_cls_token, share_embeddings)
 
     # Let's assume the model is trained
     model = WideDeep(deeptabular=deeptabular)
-    t2v = Tab2Vec(model, tab_preprocessor)
-    X_vec = t2v.transform(df_t2v)
+    t2v = Tab2Vec(model, tab_preprocessor, return_dataframe=return_dataframe)
+    t2v_out = t2v.transform(df_t2v)
 
     out_dim = (len(embed_cols) + len(cont_cols)) * deeptabular.input_dim
 
-    assert X_vec.shape[1] == out_dim
+    assert t2v_out.shape[1] == out_dim