Merge pull request #6 from flakeforever/main

Fix block slicing access in UniBlockSwap
2026-08-06 18:36:35 +08:00
parent a019e82c45 ec21b19fb2
commit 4172da4e2e
1 changed files with 6 additions and 0 deletions
@@ -131,8 +131,14 @@ class SwappableModuleList(nn.ModuleList):
        raise AttributeError(f"'{type(self).__name__}' has no attribute '{name}'")

    def __getitem__(self, idx):
+        # Support slicing: blocks[start:end]
+        if isinstance(idx, slice):
+            start, stop, step = idx.indices(self.total_count)
+            return [self[i] for i in range(start, stop, step)]
+
        if idx >= self.non_swap_count:
            self._load_swap(idx - self.non_swap_count)
+
        return super().__getitem__(idx)

    def __iter__(self):