写一个代码：基于pytorch使用softmax

import torch
import torch.nn as nn

# 定义模型
class SoftmaxModel(nn.Module):
def __init__(self, input_size, num_classes):
super(SoftmaxModel, self).__init__()
self.linear = nn.Linear(input_size, num_classes)

def forward(self, x):
out = self.linear(x)
return nn.functional.softmax(out, dim=1)

# 定义数据
input_size = 5
num_classes = 3
batch_size = 2

x = torch.randn(batch_size, input_size)

# 初始化模型
model = SoftmaxModel(input_size, num_classes)

# 前向传播
output = model(x)

print(output) # 输出softmax后的结果，每行3个数，和为1