HRM_sudoku / softmax.py
Code2aum's picture
Upload folder using huggingface_hub
5dc80b3 verified
Raw
History Blame Contribute Delete
194 Bytes
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=}")