Merge pull request #37 from MayDomine/main

add __iter__ to make TransformerBlockList Iterable
OpenBMB · Jul 8, 2022 · 3ed3e3c · 3ed3e3c
2 parents e6ba031 + bbe73a4
commit 3ed3e3c
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/bmtrain/block_layer.py b/bmtrain/block_layer.py
@@ -831,7 +831,8 @@ def __init__(self, modules: Iterable[CheckpointBlock], sqrt=False) -> None:
 
     def __len__(self) -> int:
         return len(self._modules)
-
+    def __iter__(self) -> Iterator[CheckpointBlock]:
+        return iter(self._modules.values())
     def __getitem__(self, index: Union[int, str]) -> CheckpointBlock:
         return self._modules[str(index)]