How to check if my code is currently JITed?

jit
1

I don’t use JIT for the training of my model, but only during exporting. I have shape asserts and random sampling within my code so JIT complains that:

  • assert, boolean conversion is not supported
  • the model outputs does not match over two multiple runs - because of randomness

I would like to change the behavior of my forwards depending on whether they are exported or not. Is there a method or global variable that allows me to do this kind of check? I imagine something like:

if not torch.jit.is_jit_active():
    assert x.shape == (2, 2)

I would use something similar for sampling since I’m only interested in modes of the distribution during exporting.

2

Just accidentally stumbled on the right function torch.jit.is_tracing().

My linter says it’s not exported, so I’m not sure whether it’s safe to use but it says the same for nn.Parameter so there is probably nothing to worry about. Linters never liked Pytorch anyway.