bugfix in Phi module

nn.ModuleList has to be used in order to register layer parameters as module parameters (similar to add_module)

src/nets/deepsets.py

@@ -90,13 +90,10 @@ class Phi(nn.Module):
         super(Phi, self).__init__()
         self.input_dim = input_dim
         self.output_dim = output_dim
-        self.layers = [nn.Linear(self.input_dim, hidden_dim)]
+        self.layers = nn.ModuleList([nn.Linear(self.input_dim, hidden_dim)])
         for _ in range(hidden_n - 1):
             self.layers.append(nn.Linear(hidden_dim, hidden_dim))
         self.layers.append(nn.Linear(hidden_dim, self.output_dim))
-        # self.in_layer = nn.Linear(input_dim, hidden_dim)
-        # self.hidden_layers = [nn.Linear(hidden_dim, hidden_dim) for _ in range(hidden_n - 1)]
-        # self.out_layer = nn.Linear(hidden_dim, output_dim)
         self.activation = nn.functional.relu
         self.final_activation = final_activation if final_activation else lambda x: x
 

tests/nets/test_deepsets.py

@@ -41,6 +41,8 @@ def test_phi():
     y = phi(torch.rand(input_dim))
     y = phi(torch.rand(7,7,7,input_dim))
 
+    assert len(phi.parameters() > 0)
+
 def test_deepsets():
     m = ds.DeepSetsModule.from_config(ds_config)