| 178 | self.init_method(self.tokentype_embeddings.weight) |
| 179 | |
| 180 | def forward(self, input_ids, position_ids, tokentype_ids=None): |
| 181 | # Embeddings. |
| 182 | words_embeddings = self.word_embeddings(input_ids) |
| 183 | position_embeddings = self.position_embeddings(position_ids) |
| 184 | embeddings = words_embeddings + position_embeddings |
| 185 | if tokentype_ids is not None: |
| 186 | assert self.tokentype_embeddings is not None |
| 187 | embeddings = embeddings + self.tokentype_embeddings(tokentype_ids) |
| 188 | else: |
| 189 | assert self.tokentype_embeddings is None |
| 190 | |
| 191 | # Dropout. |
| 192 | embeddings = self.embedding_dropout(embeddings) |
| 193 | |
| 194 | return embeddings |
| 195 | |
| 196 | def state_dict_for_save_checkpoint( |
| 197 | self, destination=None, prefix='', keep_vars=False, |