From ec21b19fb24d6829e778df0d5050a461ceb420ae Mon Sep 17 00:00:00 2001
From: Flakeforever <flake.chen@gmail.com>
Date: Wed, 5 Aug 2026 08:51:58 -0700
Subject: [PATCH] Fix block slicing access in UniBlockSwap

The block manager previously assumed all __getitem__ calls used integer indices. When a slice was used (e.g. blocks[start:end]), the slice object was incorrectly passed into swap logic, causing invalid comparisons and failures.
---
 block_swap.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/block_swap.py b/block_swap.py
index 746aee5..7d561a2 100644
--- a/block_swap.py
+++ b/block_swap.py
@@ -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):