From c653847fc53334f2419f130ac05a221cfe52ee77 Mon Sep 17 00:00:00 2001
From: Kashif Rasul <kashif.rasul@gmail.com>
Date: Wed, 30 Mar 2022 18:18:07 +0200
Subject: [PATCH] add param proj

---
 tft/module.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tft/module.py b/tft/module.py
index 4fa76c5..073bc82 100644
--- a/tft/module.py
+++ b/tft/module.py
@@ -508,9 +508,9 @@ class TFTModel(nn.Module):
             past_selection, future_selection, [c_h.unsqueeze(0), c_c.unsqueeze(0)]
         )
 
-        prams = self.temporal_decoder(encoding, static_enrichment)
+        dec_output = self.temporal_decoder(encoding, static_enrichment)
 
-        return prams
+        return self.param_proj(dec_output)
 
     @torch.jit.ignore
     def output_distribution(