bug fix with policy class when running on GPU

joshuaspear · Sep 6, 2024 · 1fe8136 · 1fe8136
1 parent a281468
commit 1fe8136
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/offline_rl_ope/components/Policy.py b/src/offline_rl_ope/components/Policy.py
@@ -167,7 +167,7 @@ def __call__(
         action_prs = self.postproc_tens(p_return.action_prs)
         self.collect_res_fn(action_prs)
         self.collect_act_func(actions)
-        return p_return.action_prs
+        return action_prs
 
 class GreedyDeterministic(BasePolicy):