periodically savingin out model and adding functionality to make identity Phi networks (for 0-dim NNs)

src/nets/deepsets.py

@@ -90,10 +90,14 @@ class Phi(nn.Module):
         super(Phi, self).__init__()
         self.input_dim = input_dim
         self.output_dim = output_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))
+        if hidden_n > 0:
+            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))
+        else:
+            self.layers = nn.ModuleList([nn.Identity()])
+            self.output_dim = self.input_dim
         self.activation = nn.functional.relu
         self.final_activation = final_activation if final_activation else lambda x: x