From 2bd638c2472429745929fd9ec463bc52a14de202 Mon Sep 17 00:00:00 2001 From: Johannes Fischer Date: Tue, 27 Jul 2021 12:37:40 +0200 Subject: [PATCH] Test for batch data in test_deepsets_policy --- tests/policies/test_policies.py | 20 +++++++++++++------- 1 file changed, 13 insertions(+), 7 deletions(-) diff --git a/tests/policies/test_policies.py b/tests/policies/test_policies.py index 95183e8..88a92d0 100644 --- a/tests/policies/test_policies.py +++ b/tests/policies/test_policies.py @@ -9,14 +9,23 @@ with open(config_path, 'r') as cfg: def test_deepsets_policy(): module = DeepSetsPolicy(config) + B = 25 ns = 5 nv = 7 npath = 20 - ego_state = torch.rand(ns) - relative_state = torch.rand(nv, ns) - path_x = torch.rand(npath) - path_y = torch.rand(npath) + ego_state = torch.rand(B, ns) + relative_state = torch.rand(B, nv, ns) + path_x = torch.rand(B, npath) + path_y = torch.rand(B, npath) + + sample = { + "state": ego_state[0], + "relative_state": relative_state[0], + "path_x": path_x[0], + "path_y": path_y[0], + } + module(sample) sample = { "state": ego_state, @@ -24,7 +33,4 @@ def test_deepsets_policy(): "path_x": path_x, "path_y": path_y, } - module(sample) - - # TODO: Refine test to also test for batch data \ No newline at end of file