diff --git a/tests/unit/vertexai/genai/replays/test_prompt_optimizer_optimize_job_state.py b/tests/unit/vertexai/genai/replays/test_prompt_optimizer_optimize_job_state.py
index 09bc5fc8e6..6be77a528a 100644
--- a/tests/unit/vertexai/genai/replays/test_prompt_optimizer_optimize_job_state.py
+++ b/tests/unit/vertexai/genai/replays/test_prompt_optimizer_optimize_job_state.py
@@ -44,6 +44,7 @@ def test_optimize(client):
         method="vapo",
         config=config,
     )
+    assert isinstance(job, types.CustomJob)
     assert job.state == types.JobState.JOB_STATE_SUCCEEDED
 
 
diff --git a/vertexai/_genai/prompt_optimizer.py b/vertexai/_genai/prompt_optimizer.py
index 2b742ea069..b072a9974d 100644
--- a/vertexai/_genai/prompt_optimizer.py
+++ b/vertexai/_genai/prompt_optimizer.py
@@ -618,16 +618,14 @@ def optimize(
         }
         args = ["--%s=%s" % (k, v) for k, v in container_args.items()]
         worker_pool_specs = [
-            {
-                "replica_count": 1,
-                "container_spec": {
-                    "image_uri": container_uri,
-                    "args": args,
-                },
-                "machine_spec": {
-                    "machine_type": "n1-standard-4",
-                },
-            }
+            types.WorkerPoolSpec(
+                replica_count=1,
+                machine_spec=types.MachineSpec(machine_type="n1-standard-4"),
+                container_spec=types.ContainerSpec(
+                    image_uri=container_uri,
+                    args=args,
+                ),
+            )
         ]
 
         service_account = _prompt_optimizer_utils._get_service_account(config)
@@ -879,16 +877,14 @@ async def optimize(
         }
         args = ["--%s=%s" % (k, v) for k, v in container_args.items()]
         worker_pool_specs = [
-            {
-                "replica_count": 1,
-                "container_spec": {
-                    "image_uri": container_uri,
-                    "args": args,
-                },
-                "machine_spec": {
-                    "machine_type": "n1-standard-4",
-                },
-            }
+            types.WorkerPoolSpec(
+                replica_count=1,
+                machine_spec=types.MachineSpec(machine_type="n1-standard-4"),
+                container_spec=types.ContainerSpec(
+                    image_uri=container_uri,
+                    args=args,
+                ),
+            )
         ]
 
         service_account = _prompt_optimizer_utils._get_service_account(config)