diff --git a/include/dice/dice.h b/include/dice/dice.h
index c56ce86..3565571 100644
--- a/include/dice/dice.h
+++ b/include/dice/dice.h
@@ -136,9 +136,9 @@
 // Derives a |cdi_private_key| from a |cdi_attest| value. On success populates
 // |cdi_private_key| and returns kDiceResultOk. Note: of the provided |ops|,
 // only 'kdf' is called.
-DiceResult DiceDeriveCdiPrivateKey(const DiceOps* ops,
-                                   const uint8_t cdi_attest[DICE_CDI_SIZE],
-                                   uint8_t cdi_private_key[DICE_CDI_SIZE]);
+DiceResult DiceDeriveCdiPrivateKey(
+    const DiceOps* ops, const uint8_t cdi_attest[DICE_CDI_SIZE],
+    uint8_t cdi_private_key[DICE_PRIVATE_KEY_SIZE]);
 
 // Derives an |id| from a |cdi_public_key| value. Because public keys can vary
 // in length depending on the algorithm, the |cdi_public_key_size| in bytes must
diff --git a/src/dice.c b/src/dice.c
index 0b1ff64..2c7d8a6 100644
--- a/src/dice.c
+++ b/src/dice.c
@@ -34,9 +34,9 @@
     0x1D, 0xB9, 0x52, 0x0B, 0xA5, 0x1C, 0x7B, 0x29, 0xEA};
 static const size_t kIdSaltSize = 64;
 
-DiceResult DiceDeriveCdiPrivateKey(const DiceOps* ops,
-                                   const uint8_t cdi_attest[DICE_CDI_SIZE],
-                                   uint8_t cdi_private_key[DICE_CDI_SIZE]) {
+DiceResult DiceDeriveCdiPrivateKey(
+    const DiceOps* ops, const uint8_t cdi_attest[DICE_CDI_SIZE],
+    uint8_t cdi_private_key[DICE_PRIVATE_KEY_SIZE]) {
   // Use the CDI as input key material, with fixed salt and info.
   return ops->kdf(ops, /*length=*/DICE_PRIVATE_KEY_SIZE, cdi_attest,
                   /*ikm_size=*/DICE_CDI_SIZE, kAsymSalt, kAsymSaltSize,