diff --git a/library/psa_crypto.c b/library/psa_crypto.c
index 1faf1dd..8e339ea 100644
--- a/library/psa_crypto.c
+++ b/library/psa_crypto.c
@@ -3470,7 +3470,9 @@
                                                    signature_length);
 
     /* Update ops count with work done. */
-    operation->num_ops = psa_driver_wrapper_sign_hash_get_num_ops(operation);
+    if (status == PSA_SUCCESS) {
+        status = psa_driver_wrapper_sign_hash_get_num_ops(operation, &operation->num_ops);
+    }
 
 exit:
 
@@ -3601,8 +3603,9 @@
     status = psa_driver_wrapper_verify_hash_complete(operation);
 
     /* Update ops count with work done. */
-    operation->num_ops = psa_driver_wrapper_verify_hash_get_num_ops(
-        operation);
+    if (status == PSA_SUCCESS) {
+        status = psa_driver_wrapper_verify_hash_get_num_ops(operation, &operation->num_ops);
+    }
 
 exit:
 
diff --git a/scripts/data_files/driver_templates/psa_crypto_driver_wrappers.h.jinja b/scripts/data_files/driver_templates/psa_crypto_driver_wrappers.h.jinja
index 3d116b3..4032d0e 100644
--- a/scripts/data_files/driver_templates/psa_crypto_driver_wrappers.h.jinja
+++ b/scripts/data_files/driver_templates/psa_crypto_driver_wrappers.h.jinja
@@ -472,17 +472,19 @@
     }
 }
 
-static inline uint32_t psa_driver_wrapper_sign_hash_get_num_ops(
-    psa_sign_hash_interruptible_operation_t *operation )
+static inline int psa_driver_wrapper_sign_hash_get_num_ops(
+    psa_sign_hash_interruptible_operation_t *operation, uint32_t *num_ops )
 {
     switch( operation->id )
     {
         /* If uninitialised, return 0, as no work can have been done. */
         case 0:
-            return 0;
+            *num_ops = 0;
+            return PSA_SUCCESS;
 
         case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
-            return(mbedtls_psa_sign_hash_get_num_ops(&operation->ctx.mbedtls_ctx));
+            *num_ops = mbedtls_psa_sign_hash_get_num_ops(&operation->ctx.mbedtls_ctx);
+            return PSA_SUCCESS;
 
 #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
 #if defined(PSA_CRYPTO_DRIVER_TEST)
@@ -495,17 +497,19 @@
     return( PSA_ERROR_INVALID_ARGUMENT );
 }
 
-static inline uint32_t psa_driver_wrapper_verify_hash_get_num_ops(
-    psa_verify_hash_interruptible_operation_t *operation )
+static inline int psa_driver_wrapper_verify_hash_get_num_ops(
+    psa_verify_hash_interruptible_operation_t *operation, uint32_t *num_ops )
 {
     switch( operation->id )
     {
         /* If uninitialised, return 0, as no work can have been done. */
         case 0:
-            return 0;
+            *num_ops = 0;
+            return PSA_SUCCESS;
 
         case PSA_CRYPTO_MBED_TLS_DRIVER_ID:
-            return (mbedtls_psa_verify_hash_get_num_ops(&operation->ctx.mbedtls_ctx));
+            *num_ops = mbedtls_psa_verify_hash_get_num_ops(&operation->ctx.mbedtls_ctx);
+            return PSA_SUCCESS;
 
 #if defined(PSA_CRYPTO_ACCELERATOR_DRIVER_PRESENT)
 #if defined(PSA_CRYPTO_DRIVER_TEST)
@@ -516,7 +520,7 @@
 
     }
 
-    return( PSA_ERROR_INVALID_ARGUMENT );
+    return ( PSA_ERROR_INVALID_ARGUMENT );
 }
 
 static inline psa_status_t psa_driver_wrapper_sign_hash_start(