Same as before for FakeTransforms

Author: ValerianRey

tests/unit/autojac/_transform/test_base.py

@@ -21,7 +21,7 @@ def __init__(self, required_keys: set[Tensor], output_keys: set[Tensor]):
     def __str__(self):
         return "T"
 
-    def _compute(self, input: _B) -> _C:
+    def __call__(self, input: _B) -> _C:
         # Ignore the input, create a dictionary with the right keys as an output.
         # Cast the type for the purpose of type-checking.
         output_dict = {key: torch.empty(0) for key in self._output_keys}

tests/unit/autojac/_transform/test_stack.py

@@ -17,7 +17,7 @@ class FakeGradientsTransform(Transform[EmptyTensorDict, Gradients]):
     def __init__(self, keys: Iterable[Tensor]):
         self.keys = set(keys)
 
-    def _compute(self, input: EmptyTensorDict) -> Gradients:
+    def __call__(self, input: EmptyTensorDict) -> Gradients:
         return Gradients({key: torch.ones_like(key) for key in self.keys})
 
     def check_keys(self) -> tuple[set[Tensor], set[Tensor]]: