JitTracedCache Class — pytorch Architecture

Source Code

benchmarks/dynamo/common.py lines 1469–1488

class JitTracedCache:
    cache: dict[weakref.ref, Any] = {}

    @classmethod
    def load(cls, model, example_inputs):
        key = weakref.ref(model)
        if key not in cls.cache:
            example_args, example_kwargs = _normalize_bench_inputs(example_inputs)
            if example_args:
                jit_traced_module = torch.jit.trace(
                    model, example_inputs=example_args, strict=False
                )
            else:
                jit_traced_module = torch.jit.trace(
                    model, example_kwarg_inputs=example_kwargs, strict=False
                )

            cls.cache[key] = jit_traced_module

        return cls.cache[key]