[BugFix] fix reward dim problem of shared baseline

ai4co · Oct 18, 2023 · b952d0f · b952d0f
1 parent e5f9df1
commit b952d0f
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/rl4co/models/rl/reinforce/baselines.py b/rl4co/models/rl/reinforce/baselines.py
@@ -53,8 +53,8 @@ def eval(self, td, reward, env=None):
 class SharedBaseline(REINFORCEBaseline):
     """Shared baseline: return mean of reward as baseline"""
 
-    def eval(self, td, reward, env=None, on_dim=1):  # e.g. [batch, pomo, ...]
-        return reward.mean(dim=on_dim, keepdims=True), 0
+    def eval(self, td, reward, env=None):  # e.g. [batch, pomo, ...]
+        return reward.mean(), 0
 
 
 class ExponentialBaseline(REINFORCEBaseline):