From 347e664803f0c545c85ae0b2b815bc7a452ff3db Mon Sep 17 00:00:00 2001
From: Guido1Alessandro1Trevisan
 <46954188+Guido1Alessandro1Trevisan@users.noreply.github.com>
Date: Sat, 8 Nov 2025 18:17:13 -0600
Subject: [PATCH] Remove sorting from topk in forward method

Implementation stays the same with less code and more clarity
---
 gpt_oss/torch/model.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/gpt_oss/torch/model.py b/gpt_oss/torch/model.py
index 9180d493..713082cf 100644
--- a/gpt_oss/torch/model.py
+++ b/gpt_oss/torch/model.py
@@ -312,7 +312,7 @@ def __init__(
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         t = self.norm(x)
         g = self.gate(t)
-        experts = torch.topk(g, k=self.experts_per_token, dim=-1, sorted=True)
+        experts = torch.topk(g, k=self.experts_per_token, dim=-1)
         expert_weights = torch.nn.functional.softmax(experts.values, dim=1)
         expert_indices = experts.indices