diff --git a/diced/src/sample_inputs.rs b/diced/src/sample_inputs.rs
index 92bdf76..802a6d3 100644
--- a/diced/src/sample_inputs.rs
+++ b/diced/src/sample_inputs.rs
@@ -79,9 +79,9 @@
 
     let input_values_vector = get_input_values_vector();
 
-    let (cdi_attest, cdi_seal, mut cert) = dice_ctx
-        .main_flow(UDS, UDS, &to_dice_input_values(&input_values_vector[0]))
-        .context("In make_sample_bcc_and_cdis: Trying to run first main flow.")?;
+    let (cdi_values, mut cert) =
+        dice::retry_dice_main_flow(UDS, UDS, &to_dice_input_values(&input_values_vector[0]))
+            .context("In make_sample_bcc_and_cdis: Trying to run first main flow.")?;
 
     let mut bcc: Vec<u8> = vec![];
 
@@ -93,12 +93,8 @@
     bcc.append(&mut cert);
 
     let dice_artifacts = dice::retry_bcc_main_flow(
-        &cdi_attest[..].try_into().context(
-            "In make_sample_bcc_and_cdis: Failed to convert cdi_attest to array reference. (1)",
-        )?,
-        &cdi_seal[..].try_into().context(
-            "In make_sample_bcc_and_cdis: Failed to convert cdi_seal to array reference. (1)",
-        )?,
+        &cdi_values.cdi_attest,
+        &cdi_values.cdi_seal,
         &bcc,
         &to_dice_input_values(&input_values_vector[1]),
     )