diff --git a/src/nets/deepsets.py b/src/nets/deepsets.py
index a719eee..82ba7a8 100644
--- a/src/nets/deepsets.py
+++ b/src/nets/deepsets.py
@@ -63,7 +63,7 @@ class DeepSetsModule(nn.Module):
         # shape (B, max_nv)
         notnan_mask = torch.all(~torch.isnan(x), dim=-1)
         # create zero tensor of shape (B, max_nv, latent_dim) to store phi evaluations in
-        latent = torch.zeros([*x.shape[:-1], self.latent_dim])
+        latent = torch.zeros([*x.shape[:-1], self.latent_dim], dtype=x.dtype)
         # evaluate phi for all not NaN entries
         # x[batch_dynamic_mask] has shape (notnan_mask.sum(), input_dim)
         latent[notnan_mask] = self.phi(x[notnan_mask])
diff --git a/src/policies/policy.py b/src/policies/policy.py
index 311a3f0..2ac0f22 100644
--- a/src/policies/policy.py
+++ b/src/policies/policy.py
@@ -42,6 +42,6 @@ class DeepSetsPolicy(Policy, nn.Module):
         # cat path_x, path_y to tensor of dim (B, 2*P)
         path = torch.cat([sample["path_x"], sample["path_y"]], dim=-1)
         path = self.path_net(path)
-        x = torch.cat([ego, relative, path])
+        x = torch.cat([ego, relative, path], dim=-1)
         x = self.head(x)
         return x
diff --git a/tests/policies/test_forward_pass.py b/tests/policies/test_forward_pass.py
index 9f8758b..e9be09a 100644
--- a/tests/policies/test_forward_pass.py
+++ b/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)
\ No newline at end of file
diff --git a/tests/policies/test_policies.py b/tests/policies/test_policies.py
index 7ed5153..fadc32f 100644
--- a/tests/policies/test_policies.py
+++ b/tests/policies/test_policies.py
@@ -25,4 +25,6 @@ def test_deepsets_policy():
         "path_y": path_y,
     }
     
-    module(sample)
\ No newline at end of file
+    module(sample)
+
+    # TODO: Refine test to also test for batch data
\ No newline at end of file