diff --git a/features/device_key/source/DeviceKey.cpp b/features/device_key/source/DeviceKey.cpp
index 8b3e7922ae1..fcb20724213 100644
--- a/features/device_key/source/DeviceKey.cpp
+++ b/features/device_key/source/DeviceKey.cpp
@@ -245,7 +245,7 @@ int DeviceKey::get_derived_key(uint32_t *ikey_buff, size_t ikey_size, const unsi
     return DEVICEKEY_SUCCESS;
 }
 
-int DeviceKey::generate_root_of_trust()
+int DeviceKey::generate_root_of_trust(size_t key_size)
 {
     int ret = DEVICEKEY_GENERATE_RANDOM_ERROR;
     uint32_t key_buff[DEVICE_KEY_32BYTE / sizeof(uint32_t)];
@@ -255,12 +255,16 @@ int DeviceKey::generate_root_of_trust()
         return DEVICEKEY_ALREADY_EXIST;
     }
 
+    if (key_size != DEVICE_KEY_32BYTE && key_size != DEVICE_KEY_16BYTE) {
+        return DEVICEKEY_INVALID_KEY_SIZE;
+    }
+
 #if defined(DEVICE_TRNG) || defined(MBEDTLS_ENTROPY_NV_SEED) || defined(MBEDTLS_ENTROPY_HARDWARE_ALT)
     mbedtls_entropy_context *entropy = new mbedtls_entropy_context;
     mbedtls_entropy_init(entropy);
-    memset(key_buff, 0, actual_size);
+    memset(key_buff, 0, key_size);
 
-    ret = mbedtls_entropy_func(entropy, (unsigned char *)key_buff, actual_size);
+    ret = mbedtls_entropy_func(entropy, (unsigned char *)key_buff, key_size);
     if (ret != MBED_SUCCESS) {
         ret = DEVICEKEY_GENERATE_RANDOM_ERROR;
     } else {
@@ -271,7 +275,7 @@ int DeviceKey::generate_root_of_trust()
     delete entropy;
 
     if (ret == DEVICEKEY_SUCCESS) {
-        ret = device_inject_root_of_trust(key_buff, actual_size);
+        ret = device_inject_root_of_trust(key_buff, key_size);
     }
 #endif
 
diff --git a/features/device_key/source/DeviceKey.h b/features/device_key/source/DeviceKey.h
index f32515bb325..1f595d9f593 100644
--- a/features/device_key/source/DeviceKey.h
+++ b/features/device_key/source/DeviceKey.h
@@ -110,12 +110,15 @@ class DeviceKey : private mbed::NonCopyable<DeviceKey> {
      * Uses TRNG or various other entropy sources to generate random device key and
      * inject it into device's KVStore. Device Key can only be generated once.
      *
-     * \return DEVICEKEY_SUCCESS, when device key successfully generated and injected.
-     * \return DEVICEKEY_ALREADY_EXIST, if the key has already been written.
-     * \return DEVICEKEY_GENERATE_RANDOM_ERROR if this device does not contain entropy sources and cannot generate a key.
-     * \return error codes on other failures.
+     * @param key_size Size of key in bytes to generate. Must be 16 bytes or 32 bytes. Default is 16 bytes.
+     *
+     * @return DEVICEKEY_SUCCESS, when device key successfully generated and injected.
+     * @return DEVICEKEY_ALREADY_EXIST, if the key has already been written.
+     * @return DEVICEKEY_GENERATE_RANDOM_ERROR if this device does not contain entropy sources and cannot generate a key.
+     * @return DEVICEKEY_INVALID_KEY_SIZE if key_size is not 32 or 16 bytes.
+     * @return error codes on other failures.
      */
-    int generate_root_of_trust();
+    int generate_root_of_trust(size_t key_size = DEVICE_KEY_16BYTE);
 
 private:
     // Private constructor, as class is a singleton