diff --git a/kubeflow/examples/prototypes/tf-batch-predict.jsonnet b/kubeflow/examples/prototypes/tf-batch-predict.jsonnet
index 6a85f0245b2..876e6d1216f 100644
--- a/kubeflow/examples/prototypes/tf-batch-predict.jsonnet
+++ b/kubeflow/examples/prototypes/tf-batch-predict.jsonnet
@@ -87,9 +87,7 @@ local tfBatchPredictBase = {
 
                 resources: {
                   limits: {
-                    memory: "4Gi",
-                    cpu: "4",
-                    [if base.params.numGpus > 0 then "nvidia/com/gpu"]: base.params.numGpus,
+                    [if base.params.numGpus > 0 then "nvidia.com/gpu"]: base.params.numGpus,
                   },
                 },