Train Active

2020-03-03 15:10:17 +01:00
parent 44f6589259
commit 1f612a968c
13 changed files with 102 additions and 98 deletions
--- a/lib/modules/blocks.py
+++ b/lib/modules/blocks.py
@ -106,7 +106,7 @@ class ResidualModule(nn.Module):
        self.in_shape = in_shape
        module_paramters.update(in_shape=in_shape)
        self.activation = activation() if activation else lambda x: x
-        self.residual_block = [module_class(**module_paramters) for _ in range(n)]
+        self.residual_block = nn.ModuleList([module_class(**module_paramters) for _ in range(n)])
        assert self.in_shape == self.shape, f'The in_shape: {self.in_shape} - must match the out_shape: {self.shape}.'

    def forward(self, x):