This commit is contained in:
Kashif Rasul
2022-03-30 16:55:20 +02:00
parent b805896fcf
commit 8c8b1aa4eb
+11 -2
View File
@@ -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)