wassname/pytorch-ts
1
0
Fork 0
mirror of https://github.com/wassname/pytorch-ts.git synced 2026-07-27 11:26:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

store cell_type

Browse Source
This commit is contained in:
Dr. Kashif Rasul
2020-01-13 11:53:32 +01:00
parent 2d14e9d842
commit 056f831285
1 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pts/model/tempflow/tempflow_network.py
+3 -2
View File
@@ -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,