| 409 | return hx |
| 410 | |
| 411 | def apply_op(self, input, hx): |
| 412 | fwd_mode = ( |
| 413 | BatchNorm.FwdMode.TRAINING if self.training else BatchNorm.FwdMode.INFERENCE |
| 414 | ) |
| 415 | |
| 416 | op = builtin.RNN( |
| 417 | num_layers=self.num_layers, |
| 418 | bidirectional=self.bidirectional, |
| 419 | bias=self.bias, |
| 420 | hidden_size=self.hidden_size, |
| 421 | dropout=self.dropout, |
| 422 | nonlineMode=self.nonlinearity, |
| 423 | fwd_mode=fwd_mode, |
| 424 | ) |
| 425 | output, h = apply(op, input, hx, self._flatten_weights)[:2] |
| 426 | output = output + h.sum() * 0 |
| 427 | h = h + output.sum() * 0 |
| 428 | return output, h |
| 429 | |
| 430 | |
| 431 | class LSTM(RNNBase): |