带温度的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()