Type mismatch with hooks input when using JIT

jit
4

Hey @Zodiac @caillonantoine, we need type hints to make it works properly.

from typing import Tuple

...

def hook(model, input: Tuple[torch.Tensor], output):
    pass

...
1 Like