Fix unit tests

2025-04-22 10:35:20 +02:00
parent 601b5fdbfe
commit 367d9bda7d
5 changed files with 95 additions and 75 deletions
--- a/tests/policies/test_policies.py
+++ b/tests/policies/test_policies.py
@@ -141,6 +141,7 @@ def test_policy(ds_repo_id, env_name, env_kwargs, policy_name, policy_kwargs):
    Note: We test various combinations of policy and dataset. The combinations are by no means exhaustive,
          and for now we add tests as we see fit.
    """
+    policy_kwargs["device"] = DEVICE

    train_cfg = TrainPipelineConfig(
        # TODO(rcadene, aliberts): remove dataset download