bug fix with Policy class when using GPU. Action input to policy_func…

… not being pushed to GPU.
joshuaspear · Sep 6, 2024 · a281468 · a281468
1 parent 1b1c6a2
commit a281468
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/offline_rl_ope/components/Policy.py b/src/offline_rl_ope/components/Policy.py
@@ -161,6 +161,7 @@ def __call__(
         # assert isinstance(action,torch.Tensor)
         # check_array_dim(action,2)
         state = self.preproc_tens(state)
+        action = self.preproc_tens(action)
         p_return = self.policy_func(state, action)
         actions = self.postproc_tens(p_return.actions)
         action_prs = self.postproc_tens(p_return.action_prs)