From 2443fb6eefa0e1d54251fa952012b1ce3e6822b0 Mon Sep 17 00:00:00 2001 From: Shawn Tan Date: Mon, 3 Nov 2025 17:50:29 +0000 Subject: [PATCH] Changing output to follow Transformers v5 --- scattermoe/torch-ext/scattermoe/layers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scattermoe/torch-ext/scattermoe/layers.py b/scattermoe/torch-ext/scattermoe/layers.py index b2aa535..cabe7d5 100644 --- a/scattermoe/torch-ext/scattermoe/layers.py +++ b/scattermoe/torch-ext/scattermoe/layers.py @@ -48,5 +48,5 @@ def forward(self, layer_input): gates=routing_weights ) layer_output = layer_output.view(bsz, length, emb_size) - return layer_output, router_logits + return layer_output