polish nn.Sequential attribute

Jittor · Jun 5, 2024 · 8d26bb8 · 8d26bb8
1 parent b1f18f0
commit 8d26bb8
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/python/jittor/nn.py b/python/jittor/nn.py
@@ -2308,6 +2308,15 @@ def __len__(self):
 
     def named_children(self,):
         return list(self.layers.items())
+
+    def __setattr__(self, key, value) -> None:
+        if isinstance(key, str) and key.isdigit():
+            if int(key)<len(self.layers):
+                self.add_module(key, value)
+            else:
+                super().__setattr__(key, value)
+        else:
+            super().__setattr__(key, value)
 
 
     def __getattr__(self, key):