diff --git a/pts/model/tempflow/tempflow_network.py b/pts/model/tempflow/tempflow_network.py
index 72e8b96..f4e35a5 100644
--- a/pts/model/tempflow/tempflow_network.py
+++ b/pts/model/tempflow/tempflow_network.py
@@ -19,7 +19,6 @@ class TempFlowTrainingNetwork(nn.Module):
         history_length: int,
         context_length: int,
         prediction_length: int,
-        distr_output: DistributionOutput,
         dropout_rate: float,
         lags_seq: List[int],
         target_dim: int,