softmax.py

import torch
import torch.nn.functional as F
sample = torch.tensor([[1,2,3,4,5], [5,4,3,2,1]], dtype=torch.float32, device='cuda')
ref = F.softmax(sample, dim=1)
print(f"Softmax result: {ref=}")