[Python] Fix Torch modes in LeRobot video fallback

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:

paimon-python/pypaimon/tests/multimodal_lerobot_test.py

@@ -273,6 +273,35 @@ def __call__(self, rows):
 
         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",