diff --git a/src/nets/deepsets.py b/src/nets/deepsets.py
index edbb4a0..176e72f 100644
--- a/src/nets/deepsets.py
+++ b/src/nets/deepsets.py
@@ -93,7 +93,7 @@ class Phi(nn.Module):
         # 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 self.activation
+        self.final_activation = final_activation if final_activation else lambda x: x
 
     def forward(self, x):
         for layer in self.layers[:-1]:
diff --git a/src/policies/__init__.py b/src/policies/__init__.py
new file mode 100644
index 0000000..ec9d984
--- /dev/null
+++ b/src/policies/__init__.py
@@ -0,0 +1 @@
+from src.policies.policy import DeepSetsPolicy
\ No newline at end of file