Update Wall-X to 1.1.0 (#104)

wall_x/trainer/utils/data.py

@@ -0,0 +1,46 @@
+"""Data-related utility functions used by the trainer main loop."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+
+
+def move_batch_to_device(
+    batch: Any,
+    device: torch.device,
+    *,
+    non_blocking: bool = True,
+) -> Any:
+    """Move every tensor in ``batch`` to ``device``, recursing into dict/list.
+
+    Returns a new structure with the same shape; the input ``batch`` is not
+    mutated. dict / list containers are rebuilt; tensors are moved via
+    ``.to(device, non_blocking=...)``; everything else is passed through
+    by reference.
+
+    Parameters
+    ----------
+    batch : Any
+        Typically a dict produced by the dataloader, but recursion accepts
+        dict / list / tensor / arbitrary leaves.
+    device : torch.device
+        Target device, typically ``self.device`` on the trainer.
+    non_blocking : bool
+        Whether to use pinned-memory async copy. Default True because the
+        trainer uses pinned loaders; pass False if the dataloader hasn't
+        pinned memory.
+    """
+    if isinstance(batch, dict):
+        return {
+            k: move_batch_to_device(v, device, non_blocking=non_blocking)
+            for k, v in batch.items()
+        }
+    if isinstance(batch, list):
+        return [
+            move_batch_to_device(v, device, non_blocking=non_blocking) for v in batch
+        ]
+    if isinstance(batch, torch.Tensor):
+        return batch.to(device, non_blocking=non_blocking)
+    return batch