带温度的softmax

用pytorch写一下使用带有温度的softmax的demo

import torch
import torch.nn.functional as F# 定义带有温度的softmax函数
def temperature_softmax(logits, temperature=1.0):return F.softmax(logits / temperature, dim=-1)# 输入logits
logits = torch.tensor([[1.0, 2.0, 3.0], [1.0, -1.0, 0.5]])# 不同温度的输出
temperatures = [0.5, 1.0, 2.0]for temp in temperatures:print(f"Temperature = {temp}")print(temperature_softmax(logits, temp))print()