| 80 | self.embedding_dropout = torch.nn.Dropout(cfg.hidden_dropout) |
| 81 | |
| 82 | def forward(self, input_ids): |
| 83 | # Embeddings. |
| 84 | if self.embedding_weights_in_fp32: |
| 85 | self.word_embeddings = self.word_embeddings.to(torch.float32) |
| 86 | embeddings = self.word_embeddings(input_ids) |
| 87 | if self.embedding_weights_in_fp32: |
| 88 | embeddings = embeddings.to(self.params_dtype) |
| 89 | self.word_embeddings = self.word_embeddings.to(self.params_dtype) |
| 90 | |
| 91 | # Data format change to avoid explicit transposes : [b s h] --> [s b h]. |
| 92 | embeddings = embeddings.transpose(0, 1).contiguous() |
| 93 | |
| 94 | # If the input flag for fp32 residual connection is set, convert for float. |
| 95 | if self.fp32_residual_connection: |
| 96 | embeddings = embeddings.float() |
| 97 | |
| 98 | # Dropout. |
| 99 | embeddings = self.embedding_dropout(embeddings) |
| 100 | |
| 101 | return embeddings |
| 102 | |
| 103 | |
| 104 | |