Fix wrong *g placement

RWKV · Sep 10, 2023 · d97c2b9 · d97c2b9
1 parent c19a127
commit d97c2b9
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/RWKV-v5/src/model.py b/RWKV-v5/src/model.py
@@ -326,8 +326,8 @@ def _forward_state_chunk(self, r, k, v, g, w, wk, wb, ws, x_l, last_state: TimeM
 
         # Fix missing *g for output as per :
         # https://github.com/RWKV/RWKV-infctx-trainer/commit/beb46d599042b77d53db9c7fa59a5966e7d33719#r126730367
-        return self.output(x)*g, TimeMixState(x_l, s)
-        
+        return self.output(x*g), TimeMixState(x_l, s)
+
     def _forward_chunk(self, x, last_state: TimeMixState):
         # Forward sizings (Batch, Time/ContextLength, Tokens)
         B, TT, C = x.size()