| 59 | |
| 60 | class EncoderModel(nn.Module, Seq2SeqAttrs): |
| 61 | def __init__(self, **model_kwargs): |
| 62 | nn.Module.__init__(self) |
| 63 | Seq2SeqAttrs.__init__(self, **model_kwargs) |
| 64 | self.input_dim = int(model_kwargs.get('input_dim', 1)) |
| 65 | self.seq_len = int(model_kwargs.get('seq_len')) # for the encoder |
| 66 | self.dcgru_layers = nn.ModuleList( |
| 67 | [DCGRUCell(self.rnn_units, self.max_diffusion_step, self.num_nodes, |
| 68 | filter_type=self.filter_type) for _ in range(self.num_rnn_layers)]) |
| 69 | |
| 70 | def forward(self, inputs, adj, hidden_state=None): |
| 71 | """ |