(self, sequence, dropout_prob, train=False)
| 259 | nn.init.constant_(self.decoder.linear.bias, 0) |
| 260 | |
| 261 | def sequence_embedder(self, sequence, dropout_prob, train=False) -> torch.Tensor: |
| 262 | if train: |
| 263 | batch_id_for_drop = torch.where(torch.rand(sequence.shape[0], device=sequence.device) < dropout_prob) |
| 264 | sequence[batch_id_for_drop] = 0 |
| 265 | return sequence |
| 266 | |
| 267 | |
| 268 | def forward(self, t, x, wa, wr, we, prev_x = None, prev_wa = None, train = True, **kwargs) -> torch.Tensor: |