(config, weight_dir, state_dict)
| 435 | |
| 436 | |
| 437 | def get_embedding_layer(config, weight_dir, state_dict): |
| 438 | embedding_weight = _get_embedding_weight(weight_dir, state_dict) |
| 439 | |
| 440 | model = torch.nn.Embedding(config.vocab_size, config.hidden_size, -1) |
| 441 | embed_state_dict = {} |
| 442 | embed_state_dict["weight"] = embedding_weight.to(torch.float32) |
| 443 | model.load_state_dict(embed_state_dict) |
| 444 | return model |
| 445 | |
| 446 | |
| 447 | def get_export_shapes(shapes): |