diff --git a/models/pyramid_model.py b/models/pyramid_model.py
index 2613569..0a041e8 100644
--- a/models/pyramid_model.py
+++ b/models/pyramid_model.py
@@ -154,7 +154,7 @@ class Pyramid(nn.Module):
             for param in self.body_enc.parameters(): param.requires_grad = False
             for param in self.curv_enc.parameters(): param.requires_grad = False
 
-        self.levels = [PyramidLevel({ **config['pyramid']['levels'][i] }) for i in range(len(config['pyramid']['levels']))]
+        self.levels = nn.ModuleList([PyramidLevel({ **config['pyramid']['levels'][i] }) for i in range(len(config['pyramid']['levels']))])
 
 
         # if len(config['pyramid']['resolutions']) == 2: