diff --git a/tests/nets/test_deepsets.py b/tests/nets/test_deepsets.py
index c5eb903..4abe60d 100644
--- a/tests/nets/test_deepsets.py
+++ b/tests/nets/test_deepsets.py
@@ -41,7 +41,7 @@ def test_phi():
     y = phi(torch.rand(input_dim))
     y = phi(torch.rand(7,7,7,input_dim))
 
-    assert len(phi.parameters() > 0)
+    assert len(list(phi.parameters())) > 0
 
 def test_deepsets():
     m = ds.DeepSetsModule.from_config(ds_config)
diff --git a/tests/policies/test_forward_pass.py b/tests/policies/test_forward_pass.py
index 2bd314f..822688a 100644
--- a/tests/policies/test_forward_pass.py
+++ b/tests/policies/test_forward_pass.py
@@ -11,11 +11,11 @@ filestr = 'tests/policies/base_'
 
 
 def test_forward_pass():
-batch = pickle.load(open(filestr + '_test_batch.pkl', 'rb'))
-batch['relative_state'] = batch['relative_state'].float()
-policy = BehaviorCloningPolicy.load_model(config, filestr)
-policy.policy = policy.policy.type(batch['state'].dtype)
-action = policy(batch)
-i = torch.where(action.isnan())[0]
-for key in batch.keys():
-    print(batch[key][i])
+    batch = pickle.load(open(filestr + '_test_batch.pkl', 'rb'))
+    batch['relative_state'] = batch['relative_state'].float()
+    policy = BehaviorCloningPolicy.load_model(config, filestr)
+    policy.policy = policy.policy.type(batch['state'].dtype)
+    action = policy(batch)
+    i = torch.where(action.isnan())[0]
+    for key in batch.keys():
+        print(batch[key][i])
diff --git a/tests/policies/test_policies.py b/tests/policies/test_policies.py
index fadc32f..95183e8 100644
--- a/tests/policies/test_policies.py
+++ b/tests/policies/test_policies.py
@@ -7,7 +7,7 @@ with open(config_path, 'r') as cfg:
     config = json5.load(cfg)
 
 def test_deepsets_policy():
-    module = DeepSetsPolicy(config["ego_state"], config["deepsets"], config["path_encoder"], config["head"])
+    module = DeepSetsPolicy(config)
     
     ns = 5 
     nv = 7
diff --git a/tests/test.py b/tests/test.py
deleted file mode 100644
index e69de29..0000000