reset notensordict test

pytorch · Oct 3, 2023 · a24ab8d · a24ab8d
1 parent 8d36787
commit a24ab8d
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/test/test_cost.py b/test/test_cost.py
@@ -2419,8 +2419,10 @@ def test_td3_notensordict(
             loss_val = loss(**kwargs)
             for i in loss_val:
                 assert i in loss_val_td.values(), f"{i} not in {loss_val_td.values()}"
-            # for i, key in enumerate(loss_val_td.keys()):
-            # torch.testing.assert_close(loss_val_td.get(key), loss_val[i])
+
+            for i, key in enumerate(loss.out_keys):
+                torch.testing.assert_close(loss_val_td.get(key), loss_val[i])
+
             # test select
             loss.select_out_keys("loss_actor", "loss_qvalue")
             torch.manual_seed(0)