(self, input_)
| 109 | self.num_embeddings_per_partition, 0, init_method) |
| 110 | |
| 111 | def forward(self, input_): |
| 112 | # Build the mask. |
| 113 | input_mask = (input_ < self.vocab_start_index) | \ |
| 114 | (input_ >= self.vocab_end_index) |
| 115 | # Mask the input. |
| 116 | masked_input = input_.clone() - self.vocab_start_index |
| 117 | masked_input[input_mask] = 0 |
| 118 | # Get the embeddings. |
| 119 | output_parallel = F.embedding(masked_input, self.weight, |
| 120 | self.padding_idx, self.max_norm, |
| 121 | self.norm_type, self.scale_grad_by_freq, |
| 122 | self.sparse) |
| 123 | # Mask the output embedding. |
| 124 | output_parallel[input_mask, :] = 0.0 |
| 125 | # Reduce across all the model parallel GPUs. |
| 126 | output = reduce_from_model_parallel_region(output_parallel) |
| 127 | return output |
| 128 | |
| 129 | |
| 130 | class ParallelEmbedding(torch.nn.Module): |
nothing calls this directly
no test coverage detected