def embed_text(model, tokens, attention_mask, token_types):
        """Embed a batch of tokens using the model"""
        logits = model(tokens,
                              attention_mask,
                              token_types)
        return logits
Embed a batch of tokens using the model
training/DeepSpeed-Domino/megatron/model/biencoder_model.py