feat(models): vit uses patched sequential

todd/models/modules/vit.py

@@ -7,6 +7,7 @@
 from torch import nn
 
 from ...bases.configs import Config
+from ...patches.torch import Sequential
 from ...registries import InitWeightsMixin
 from ..utils import interpolate_position_embedding
 
@@ -88,7 +89,7 @@ def __init__(
         self._position_embedding = nn.Parameter(
             torch.empty(self.num_patches + 1, width),
         )
-        self._blocks = nn.Sequential(
+        self._blocks = Sequential(
             *[
                 self.BLOCK_TYPE(width=width, num_heads=num_heads)
                 for _ in range(depth)

todd/patches/torch/nn.py

@@ -33,9 +33,17 @@ def forward(self, *args, **kwargs) -> dict[str, nn.Module]:
 
 class Sequential(nn.Sequential):
 
+    def __init__(self, *args, unpack_args: bool = False, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self._unpack_args = unpack_args
+
     def forward(self, *args, **kwargs) -> tuple[Any, ...]:
+        if not self._unpack_args:
+            args, = args
         for m in self:
-            args = m(*args, **kwargs)
+            args = (
+                m(*args, **kwargs) if self._unpack_args else m(args, **kwargs)
+            )
         return args