---
title: Embedding
description: "将离散ID映射为稠密向量，目前KuDNN支持torch.float16数据类型，其他数据类型会走开源分支。"
url: https://www.hikunpeng.com/document/detail/zh/kunpenghpcs/hpckit/devg/KunpengHPCKit_developer_209.html
sourcePath: /source/zh/kunpenghpcs/hpckit/devg/KunpengHPCKit_developer_209.html
indexId: 7a9ea6785bc165b1a18e3ea606dcb7fe9986e10f304673554da665a25e014a0767
---
# Embedding

#### 场景说明

将离散ID映射为稠密向量，目前KuDNN支持torch.float16数据类型，其他数据类型会走开源分支。


#### 示例代码

```
import torch
import torch.nn as nn

# 启用kdnn
torch._C._set_kdnn_enabled(True)

# 构造Embedding层
embed = nn.Embedding(num_embeddings=1000, embedding_dim=128, dtype=torch.float16)

# 输入数据2x2的token索引
input_ids = torch.LongTensor([[1, 2], [3, 4]])
embeddings = embed(input_ids)  # 输出 [2, 2, 128]

print(embeddings)
```