diff --git a/tests/integration_tests/test_vllm_policy_correctness.py b/tests/integration_tests/test_vllm_policy_correctness.py
index e2da9b068..71ff3677b 100644
--- a/tests/integration_tests/test_vllm_policy_correctness.py
+++ b/tests/integration_tests/test_vllm_policy_correctness.py
@@ -6,7 +6,7 @@
 
 import pytest
 
-from forge.actors.policy import Policy
+from forge.actors.generator import Generator as Policy
 from vllm import SamplingParams
 from vllm.engine.arg_utils import AsyncEngineArgs
 from vllm.sampling_params import RequestOutputKind