diff --git a/src/gail/options.py b/src/gail/options.py
index 21dd7af..8eb5f16 100644
--- a/src/gail/options.py
+++ b/src/gail/options.py
@@ -36,7 +36,7 @@ class OptionsEnv(gym.Wrapper):
                 self.episode_start = True
 
             self.m = available_actions(self.env, self.options)
-            self.ch, self.value, self.log_prob = generator.policy.predict({
+            self.ch, self.value, self.log_prob = generator.policy.forward({
                 'obs': torch.tensor(self.s).unsqueeze(0).to(generator.policy.device),
                 'mask': torch.tensor(self.m).unsqueeze(0).to(generator.policy.device),
             })
diff --git a/src/policies/options.py b/src/policies/options.py
index 9bf71c9..851edd4 100644
--- a/src/policies/options.py
+++ b/src/policies/options.py
@@ -22,7 +22,7 @@ class OptionsCnnPolicy(stable_baselines3.common.policies.ActorCriticCnnPolicy):
         values = self.value_net(latent_vf)
         return values, distribution.distribution
 
-    def predict(self, obs, eps=1e-6):
+    def forward(self, obs, eps=1e-6):
         """
         Will mask invalid states before making action selections
         Args: