pyramid model fix

2023-05-22 23:20:15 +02:00
parent 08bf60a76a
commit 7d17b1a222
1 changed files with 1 additions and 1 deletions
@@ -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: