This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new 8bd0f7f19c [Python] Fix Torch modes in LeRobot video fallback (#10731)
8bd0f7f19c is described below
commit 8bd0f7f19ca8c8a21f9948cbe26e0ad8c3a24a37
Author: XiaoHongbo <[email protected]>
AuthorDate: Thu Oct 22 09:06:24 2026 +0800
[Python] Fix Torch modes in LeRobot video fallback (#10731)
---
.../pypaimon/multimodal/lerobot/dataset.py | 14 ++++++++++-
.../pypaimon/tests/multimodal_lerobot_test.py | 29 ++++++++++++++++++++++
2 files changed, 42 insertions(+), 1 deletion(-)
diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py
b/paimon-python/pypaimon/multimodal/lerobot/dataset.py
index 1ce3998ae4..5456ff7cf7 100644
--- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py
+++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py
@@ -2267,12 +2267,24 @@ def _decode_video_rows(
else:
from concurrent.futures import wait
+ torch = sys.modules.get("torch")
+ if torch is not None:
+ grad_enabled = torch.is_grad_enabled()
+ inference_enabled = torch.is_inference_mode_enabled()
+
+ def decode(collator, rows):
+ if torch is None:
+ return collator(rows)
+ with torch.inference_mode(inference_enabled), \
+ torch.set_grad_enabled(grad_enabled):
+ return collator(rows)
+
pool = decode_pool if decode_pool is not None else _VideoDecodePool()
futures = []
try:
executor = pool.get(decode_parallelism)
for collator, _, rows in tasks:
- futures.append(executor.submit(collator, rows))
+ futures.append(executor.submit(decode, collator, rows))
decoded_groups = [future.result() for future in futures]
finally:
for future in futures:
diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py
b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py
index acd87cccb0..7d48264a3a 100644
--- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py
+++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py
@@ -273,6 +273,35 @@ class LeRobotValidationTest(unittest.TestCase):
self.assertEqual([8], worker_counts)
+ def test_video_row_fallback_preserves_torch_mode(self):
+ try:
+ import torch
+ except ImportError as error:
+ self.skipTest(str(error))
+ from pypaimon.multimodal.video import _VideoDecodePool
+
+ class Collator:
+ video_column = "camera"
+ output_column = "frame"
+
+ def __call__(self, rows):
+ mode = (torch.is_grad_enabled(),
+ torch.is_inference_mode_enabled())
+ return [dict(row, frame=mode) for row in rows]
+
+ pool = _VideoDecodePool()
+ try:
+ for context, expected in (
+ (torch.enable_grad(), (True, False)),
+ (torch.no_grad(), (False, False)),
+ (torch.inference_mode(), (False, True))):
+ with self.subTest(expected=expected), context:
+ rows = {0: {"camera": 1}}
+ _decode_video_rows([rows], [Collator()], decode_pool=pool)
+ self.assertEqual(expected, rows[0]["frame"])
+ finally:
+ pool.close()
+
@unittest.skipUnless(
av is not None and importlib.util.find_spec("torch") is not None,
"PyAV and Torch are required for video decoding",
