| 109 | |
| 110 | class LSTM(torch.jit.ScriptModule): |
| 111 | def __init__(self, input_size, hidden_size, bidirectional=False, dropout=0., cell_factory=RecurrentDropoutLSTMCell): |
| 112 | super(LSTM, self).__init__() |
| 113 | self.input_size = input_size |
| 114 | self.hidden_size = hidden_size |
| 115 | self.bidirectional = bidirectional |
| 116 | self.dropout = dropout |
| 117 | self.cell_factory = cell_factory |
| 118 | num_directions = 2 if bidirectional else 1 |
| 119 | self.lstm_cells = [] |
| 120 | |
| 121 | for direction in range(num_directions): |
| 122 | cell = cell_factory(input_size, hidden_size, dropout=dropout) |
| 123 | self.lstm_cells.append(cell) |
| 124 | |
| 125 | suffix = '_reverse' if direction == 1 else '' |
| 126 | cell_name = 'cell{}'.format(suffix) |
| 127 | self.add_module(cell_name, cell) |
| 128 | |
| 129 | def forward(self, input, hidden_state=None): |
| 130 | is_packed = isinstance(input, torch.nn.utils.rnn.PackedSequence) |