From 056f8312854554af845bd8a0bc0dcd9f2411a799 Mon Sep 17 00:00:00 2001 From: "Dr. Kashif Rasul" Date: Mon, 13 Jan 2020 11:53:32 +0100 Subject: [PATCH] store cell_type --- pts/model/tempflow/tempflow_network.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/pts/model/tempflow/tempflow_network.py b/pts/model/tempflow/tempflow_network.py index f4e35a5..7b94d9a 100644 --- a/pts/model/tempflow/tempflow_network.py +++ b/pts/model/tempflow/tempflow_network.py @@ -38,7 +38,8 @@ class TempFlowTrainingNetwork(nn.Module): lags_seq.sort() self.lags_seq = lags_seq - rnn_cls = {"LSTM": nn.LSTM, "GRU": nn.GRU}[self.cell_type] + self.cell_type = cell_type + rnn_cls = {"LSTM": nn.LSTM, "GRU": nn.GRU}[cell_type] self.rnn = rnn_cls( input_size=input_size, hidden_size=num_cells, @@ -47,7 +48,7 @@ class TempFlowTrainingNetwork(nn.Module): batch_first=True, ) - flow_cls = {"RealNVP": RealNVP}[self.flow_type] + flow_cls = {"RealNVP": RealNVP}[flow_type] self.flow = flow_cls( input_size=input_size, n_blocks=n_blocks,