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.