diff --git a/fms_extras/utils/cache/paged.py b/fms_extras/utils/cache/paged.py index a3235c7..37cf417 100644 --- a/fms_extras/utils/cache/paged.py +++ b/fms_extras/utils/cache/paged.py @@ -258,14 +258,12 @@ def has_side_effects(self): @classmethod def create(cls, kernel, *args, mutated_inputs=[], **kwargs) -> None: with V.graph.fake_mode: - ( - example_output, - tensor_args, - non_tensor_args, - unflatten_args, - ) = cls.process_kernel( - kernel, *args, **kwargs - ) # type: ignore + # process_kernel returns 5 values since PyTorch 2.x added + # unbacked_bindings as the fifth element; unpack with * to stay + # compatible across versions. + example_output, tensor_args, non_tensor_args, unflatten_args, *_ = ( + cls.process_kernel(kernel, *args, **kwargs) + ) for tensor_arg in tensor_args: tensor_arg.realize()