File size: 194 Bytes
5dc80b3 | 1 2 3 4 5 | 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=}") |
5dc80b3 | 1 2 3 4 5 | 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=}") |