diff --git a/tft/lightning_module.py b/tft/lightning_module.py
index a346427..3a3465f 100644
--- a/tft/lightning_module.py
+++ b/tft/lightning_module.py
@@ -59,7 +59,13 @@ class TFTLightningModule(pl.LightningModule):
         future_target = batch["future_target"]
         future_observed_values = batch["future_observed_values"]
 
-        tft_target, time_feat, scale, static_feat = self.model.create_network_inputs(
+        (
+            tft_target,
+            time_feat,
+            scale,
+            embedded_cat,
+            static_feat,
+        ) = self.model.create_network_inputs(
             feat_static_cat=feat_static_cat,
             feat_static_real=feat_static_real,
             past_time_feat=past_time_feat,
@@ -69,7 +75,10 @@ class TFTLightningModule(pl.LightningModule):
             future_target=future_target,
         )
         params = self.model.output_params(
-            tft_target, time_feat=time_feat, static_feat=static_feat
+            tft_target,
+            time_feat=time_feat,
+            embedded_cat=embedded_cat,
+            static_feat=static_feat,
         )
         distr = self.model.output_distribution(params, scale)