tests for tdmpc and diffusion policy are passing

2024-04-09 02:50:32 +00:00
parent 1e09507bc1
commit 73dfa3c8e3
6 changed files with 19 additions and 17 deletions
--- a/tests/test_policies.py
+++ b/tests/test_policies.py
@@ -49,7 +49,7 @@ def test_policy(env_name, policy_name, extra_overrides):
    dataloader = torch.utils.data.DataLoader(
        dataset,
        num_workers=4,
-        batch_size=cfg.policy.batch_size,
+        batch_size=2,
        shuffle=True,
        pin_memory=DEVICE != "cpu",
        drop_last=True,