Resolve #72: Make self._layers a nn.ModuleList variable (#74)

asyml · Jun 26, 2019 · 57ad85d · 57ad85d
1 parent e72a966
commit 57ad85d
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/texar/core/layers.py b/texar/core/layers.py
@@ -635,12 +635,12 @@ def __init__(self, layers: Optional[List[nn.Module]] = None,
         self._mode = mode
         self._dim = dim
 
-        self._layers: Optional[List[nn.Module]] = None
+        self._layers: Optional[nn.ModuleList] = None
         if layers is not None:
             if len(layers) == 0:
                 raise ValueError(
                     "'layers' must be either None or a non-empty list.")
-            self._layers = []
+            self._layers = nn.ModuleList()
             for layer in layers:
                 if isinstance(layer, nn.Module):
                     self._layers.append(layer)
@@ -704,7 +704,7 @@ def forward(self, input: torch.Tensor) -> torch.Tensor:  # type: ignore
         return outputs
 
     @property
-    def layers(self) -> Optional[List[nn.Module]]:
+    def layers(self) -> Optional[nn.ModuleList]:
         r"""The list of parallel layers.
         """
         return self._layers