diff --git a/plugins/flytekit-kf-pytorch/tests/test_pytorch_task.py b/plugins/flytekit-kf-pytorch/tests/test_pytorch_task.py
index 7e81c189bf3..1e796fea453 100644
--- a/plugins/flytekit-kf-pytorch/tests/test_pytorch_task.py
+++ b/plugins/flytekit-kf-pytorch/tests/test_pytorch_task.py
@@ -5,7 +5,7 @@
 
 
 def test_pytorch_task():
-    @task(task_config=PyTorch(num_workers=10, per_replica_requests=Resources(cpu="1")), cache=True, cache_version="1")
+    @task(task_config=PyTorch(num_workers=10), cache=True, cache_version="1", requests=Resources(cpu="1"),)
     def my_pytorch_task(x: int, y: str) -> int:
         return x
 
diff --git a/plugins/flytekit-kf-tensorflow/tests/test_tensorflow_task.py b/plugins/flytekit-kf-tensorflow/tests/test_tensorflow_task.py
index 038f0856ed3..6e3b5474cd5 100644
--- a/plugins/flytekit-kf-tensorflow/tests/test_tensorflow_task.py
+++ b/plugins/flytekit-kf-tensorflow/tests/test_tensorflow_task.py
@@ -7,9 +7,10 @@
 def test_tensorflow_task():
     @task(
         task_config=TfJob(
-            num_workers=10, per_replica_requests=Resources(cpu="1"), num_ps_replicas=1, num_chief_replicas=1
+            num_workers=10,  num_ps_replicas=1, num_chief_replicas=1
         ),
         cache=True,
+        requests=Resources(cpu="1"),
         cache_version="1",
     )
     def my_tensorflow_task(x: int, y: str) -> int: