add check_keys in test_stack.FakeGradientsTransform

Author: PierreQuinton

tests/unit/autojac/_transform/test_stack.py

@@ -20,13 +20,8 @@ def __init__(self, keys: Iterable[Tensor]):
     def _compute(self, input: EmptyTensorDict) -> Gradients:
         return Gradients({key: torch.ones_like(key) for key in self.keys})
 
-    @property
-    def required_keys(self) -> set[Tensor]:
-        return set()
-
-    @property
-    def output_keys(self) -> set[Tensor]:
-        return self.keys
+    def check_keys(self) -> tuple[set[Tensor], set[Tensor]]:
+        return set(), self.keys
 
 
 def test_single_key():