Keep compiled cache copies on the compute stream

Author: aw471

sdm/models/base/model.py

@@ -50,7 +50,8 @@ def _prefetch(
     ) -> "_Prefetcher[T]":
 
         stream: torch.cuda.Stream | None = None
-        if device.type == "cuda":
+        # Keep compiled copies and compute on one stream for safe buffer reuse.
+        if device.type == "cuda" and not torch.compiler.is_compiling():
             if device not in self._transfer_streams:
                 self._transfer_streams[device] = torch.cuda.Stream(device)
             stream = self._transfer_streams[device]
@@ -121,7 +122,10 @@ def __enter__(self) -> Self:
                 self._transfer_stream.synchronize()
                 raise
         else:
-            self._next = self._items[0].to(self._device)  # type: ignore
+            self._next = self._items[0].to(  # type: ignore
+                device=self._device,
+                non_blocking=self._device.type == "cuda",
+            )
 
         return self
 
@@ -159,6 +163,9 @@ def __iter__(self) -> Iterator[T]:
                             non_blocking=True,
                         )
                 else:
-                    self._next = self._next.to(self._device)  # type: ignore
+                    self._next = self._next.to(  # type: ignore
+                        device=self._device,
+                        non_blocking=self._device.type == "cuda",
+                    )
 
             yield item