Clarify Composition.check_keys

Author: ValerianRey

src/torchjd/autojac/_transform/base.py

@@ -80,7 +80,8 @@ def __call__(self, input: _A) -> _C:
 
     def check_keys(self, input_keys: set[Tensor]) -> set[Tensor]:
         intermediate_keys = self.inner.check_keys(input_keys)
-        return self.outer.check_keys(intermediate_keys)
+        output_keys = self.outer.check_keys(intermediate_keys)
+        return output_keys
 
 
 class Conjunction(Transform[_A, _B]):