diff --git a/fault_injection/fi_otbn.py b/fault_injection/fi_otbn.py
index 8c907cc9..47449514 100755
--- a/fault_injection/fi_otbn.py
+++ b/fault_injection/fi_otbn.py
@@ -101,10 +101,10 @@ def fi_parameter_sweep(cfg: dict, target: Target, fi_gear,
     Returns:
         device_id: The ID of the target device.
     """
-    # Setup key manager if needed by test.
-    ot_communication.init_keymgr(cfg["test"]["which_test"])
     # Configure the OTBN FI code on the target.
     device_id = ot_communication.init()
+    # Setup key manager if needed by test.
+    ot_communication.init_keymgr(cfg["test"]["which_test"])
     # Store results in array for a quick access.
     fi_results = []
     # Start the parameter sweep.
@@ -132,10 +132,10 @@ def fi_parameter_sweep(cfg: dict, target: Target, fi_gear,
                 ot_communication = target.reset_target(com_reset = True)
                 # Re-establish UART connection.
                 ot_communication = OTFIOtbn(target)
-                # Setup key manager if needed by test.
-                ot_communication.init_keymgr(cfg["test"]["which_test"])
                 # Configure the OTBN FI code on the target.
                 ot_communication.init()
+                # Setup key manager if needed by test.
+                ot_communication.init_keymgr(cfg["test"]["which_test"])
                 # Reset FIGear if necessary.
                 fi_gear.reset()
             else:
diff --git a/target/communication/fi_otbn_commands.py b/target/communication/fi_otbn_commands.py
index 371c2b35..fdf22c54 100644
--- a/target/communication/fi_otbn_commands.py
+++ b/target/communication/fi_otbn_commands.py
@@ -84,13 +84,7 @@ def init_keymgr(self, test: str) -> None:
             # InitTrigger command.
             time.sleep(0.01)
             self.target.write(json.dumps("InitKeyMgr").encode("ascii"))
-            # As the init resets the chip, we need to call it again to complete
-            # the initialization of the key manager.
-            time.sleep(2)
-            self._ujson_otbn_fi_cmd()
-            time.sleep(0.01)
-            self.target.write(json.dumps("InitKeyMgr").encode("ascii"))
-            time.sleep(2)
+            time.sleep(0.5)
 
     def init(self) -> None:
         """ Initialize the OTBN FI code on the chip.