bugfix in deepsets

tests/policies/test_forward_pass.py

@@ -9,6 +9,9 @@ with open(config_path, 'r') as cfg:
 
 filestr = 'tests/policies/base_'
 
+
 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)

tests/policies/test_policies.py

@@ -25,4 +25,6 @@ def test_deepsets_policy():
         "path_y": path_y,
     }
     
-    module(sample)
+    module(sample)
+
+    # TODO: Refine test to also test for batch data