fix: batch transform

Author: mjun0812

README.md

@@ -10,7 +10,7 @@ A comprehensive, production-ready PyTorch project template with modular architec
 - **📊 Experiment Tracking**: MLflow and Weights & Biases integration with auto-visualization
 - **🔧 Modern Tooling**: uv package management, pre-commit hooks, Docker support
 - **💾 Resume Training**: Automatic checkpoint saving and loading with state preservation
-- **🌐 Cross-Platform**: Development support on macOS and Linux with optimized builds
+- **🌐 Cross-Platform**: Development support on macOS (Apple Silicon MPS), Linux with optimized builds
 - **🐳 Development Environment**: Devcontainer and Jupyter Lab integration
 - **⚡ Performance Optimization**: RAM caching, mixed precision, torch.compile support
 - **📚 Auto Documentation**: Sphinx-based API docs with live reloading

src/transform/__init__.py

@@ -5,6 +5,7 @@
 from .build import (
     BATCHED_TRANSFORM_REGISTRY,
     TRANSFORM_REGISTRY,
+    BatchedTransformCompose,
     build_batched_transform,
     build_transform,
     build_transforms,

src/transform/batch_compose.py

@@ -0,0 +1,28 @@
+import torch
+import torch.nn as nn
+
+from ..dataloaders import DatasetOutput
+
+
+class BatchedTransformCompose:
+    def __init__(self, transforms: list[nn.Module]) -> None:
+        self.transforms = transforms
+
+    def to(self, device: torch.device) -> "BatchedTransformCompose":
+        for t in self.transforms:
+            if isinstance(t, torch.nn.Module):
+                t.to(device)
+        return self
+
+    def __call__(self, data: DatasetOutput) -> DatasetOutput:
+        for t in self.transforms:
+            data = t(data)
+        return data
+
+    def __repr__(self) -> str:
+        format_string = self.__class__.__name__ + "("
+        for t in self.transforms:
+            format_string += "\n"
+            format_string += f"    {t}"
+        format_string += "\n)"
+        return format_string

src/transform/build.py

@@ -3,6 +3,7 @@
 from ..config import TransformConfig
 from ..utils import Registry
 from .base import BaseTransform
+from .batch_compose import BatchedTransformCompose
 
 TRANSFORM_REGISTRY = Registry("TRANSFORM")
 """Registry for data transformation classes."""
@@ -54,7 +55,7 @@ def build_transforms(cfg: list[TransformConfig]) -> T.Compose:
     return T.Compose(transforms)
 
 
-def build_batched_transform(cfg: list[TransformConfig]) -> T.Compose:
+def build_batched_transform(cfg: list[TransformConfig]) -> BatchedTransformCompose:
     """Build a composition of batch-level transformations.
 
     Args:
@@ -66,4 +67,4 @@ def build_batched_transform(cfg: list[TransformConfig]) -> T.Compose:
     batched_transforms = []
     for cfg_transform in cfg:
         batched_transforms.append(build_batch_transform(cfg_transform))
-    return T.Compose(batched_transforms)
+    return BatchedTransformCompose(batched_transforms)