From 29f6abc997fd0fe6f4de685c33a36110603cebb1 Mon Sep 17 00:00:00 2001
From: Remi Cadene <re.cadene@gmail.com>
Date: Tue, 15 Oct 2024 19:38:40 +0200
Subject: [PATCH] Add safe_disconnect to replay

---
 lerobot/scripts/control_robot.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/lerobot/scripts/control_robot.py b/lerobot/scripts/control_robot.py
index 3fdef1ff5..425247e65 100644
--- a/lerobot/scripts/control_robot.py
+++ b/lerobot/scripts/control_robot.py
@@ -310,6 +310,7 @@ def record(
     return lerobot_dataset
 
 
+@safe_disconnect
 def replay(
     robot: Robot, episode: int, fps: int | None = None, root="data", repo_id="lerobot/debug", play_sounds=True
 ):