---
title: SoftMax
description: "多分类任务中沿指定维度计算归一化指数概率分布，目前KuDNN支持torch.float32数据类型，其他数据类型会使用开源分支。"
url: https://www.hikunpeng.com/document/detail/zh/kunpenghpcs/hpckit/devg/KunpengHPCKit_developer_204.html
sourcePath: /source/zh/kunpenghpcs/hpckit/devg/KunpengHPCKit_developer_204.html
indexId: b5a42a2cf8cf24af2a84de7fab6be4334d3ef801040d4b68c32e9b0dfd18698867
---
# SoftMax

#### 场景说明

多分类任务中沿指定维度计算归一化指数概率分布，目前KuDNN支持torch.float32数据类型，其他数据类型会使用开源分支。


#### 示例代码

```
import torch
import torch.nn as nn

#使能KuDNN
torch._C._set_kdnn_enabled(True)

# 构造Softmax层对列进行归一化
m = nn.Softmax(dim=1)
logits = torch.randn(4, 5)  # 4个样本，5分类
probs = m(logits)

# 每行概率和为1
print(probs)
```