minor simplify

This commit is contained in:
GeekExplorer
2026-04-13 22:09:46 +08:00
parent 02a95fdc66
commit 9e8507ef41
6 changed files with 14 additions and 27 deletions
-3
View File
@@ -4,9 +4,6 @@ from torch import nn
class Sampler(nn.Module):
def __init__(self):
super().__init__()
@torch.compile
def forward(self, logits: torch.Tensor, temperatures: torch.Tensor):
logits = logits.float().div_(temperatures.unsqueeze(dim=1))