Merge pull request #22 from DeepTrackAI/bm/allow-non-deeplay-list-get

Bm/allow-non-deeplay-list-get
DeepTrackAI · Dec 4, 2023 · eda2d02 · eda2d02
2 parents a1590fa + f2f67ed
commit eda2d02
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 2 deletions.
diff --git a/deeplay/list.py b/deeplay/list.py
@@ -89,8 +89,7 @@ def __getattr__(self, name: str) -> "LayerList[T]":
             submodules = [
                 getattr(layer, name)
                 for layer in self
-                if hasattr(layer, name)
-                and isinstance(getattr(layer, name), DeeplayModule)
+                if hasattr(layer, name) and isinstance(getattr(layer, name), nn.Module)
             ]
             if len(submodules) > 0:
                 return LayerList(*submodules)

diff --git a/deeplay/tests/test_layerlist.py b/deeplay/tests/test_layerlist.py
@@ -165,6 +165,24 @@ def test_with_instantiated(self):
         self.assertIsInstance(llist[0], nn.Linear)
         self.assertIsInstance(llist[1], nn.Linear)
 
+    def test_with_instantiated_2(self):
+        class Item(DeeplayModule):
+            def __init__(self):
+                self.net = Layer(nn.Linear, 1, 1)
+
+        llist = LayerList(Item(), Item())
+
+        nets = llist.net
+        self.assertEqual(len(nets), 2)
+        self.assertIsInstance(nets[0], Layer)
+        self.assertIsInstance(nets[1], Layer)
+
+        llist.build()
+        nets = llist.net
+        self.assertEqual(len(nets), 2)
+        self.assertIsInstance(nets[0], nn.Linear)
+        self.assertIsInstance(nets[1], nn.Linear) 
+
     def test_slice_does_not_mutate(self):
         llist = LayerList(
             LayerActivation(Layer(nn.Linear, 1, 1), Layer(nn.ReLU)),
@@ -178,3 +196,4 @@ def test_slice_does_not_mutate(self):
             self.assertTrue(layer._has_built)
         for layer in llist[0:]:
             self.assertTrue(layer._has_built)
+