diff --git a/programs/psa/crypto_examples.c b/programs/psa/crypto_examples.c
index 07d1fd2..72fa12f 100644
--- a/programs/psa/crypto_examples.c
+++ b/programs/psa/crypto_examples.c
@@ -162,9 +162,9 @@
                              PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT );
     psa_set_key_algorithm( &attributes, alg );
     psa_set_key_type( &attributes, PSA_KEY_TYPE_AES );
+    psa_set_key_bits( &attributes, key_bits );
 
-    status = psa_generate_key( &attributes, &key_handle, key_bits,
-                               NULL, 0 );
+    status = psa_generate_key( &attributes, &key_handle, NULL, 0 );
     ASSERT_STATUS( status, PSA_SUCCESS );
 
     status = cipher_encrypt( key_handle, alg, iv, sizeof( iv ),
@@ -213,9 +213,9 @@
                              PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT );
     psa_set_key_algorithm( &attributes, alg );
     psa_set_key_type( &attributes, PSA_KEY_TYPE_AES );
+    psa_set_key_bits( &attributes, key_bits );
 
-    status = psa_generate_key( &attributes, &key_handle, key_bits,
-                               NULL, 0 );
+    status = psa_generate_key( &attributes, &key_handle, NULL, 0 );
     ASSERT_STATUS( status, PSA_SUCCESS );
 
     status = cipher_encrypt( key_handle, alg, iv, sizeof( iv ),
@@ -260,9 +260,9 @@
                              PSA_KEY_USAGE_ENCRYPT | PSA_KEY_USAGE_DECRYPT );
     psa_set_key_algorithm( &attributes, alg );
     psa_set_key_type( &attributes, PSA_KEY_TYPE_AES );
+    psa_set_key_bits( &attributes, key_bits );
 
-    status = psa_generate_key( &attributes, &key_handle, key_bits,
-                               NULL, 0 );
+    status = psa_generate_key( &attributes, &key_handle, NULL, 0 );
     ASSERT_STATUS( status, PSA_SUCCESS );
 
     status = cipher_encrypt( key_handle, alg, iv, sizeof( iv ),