Find if a module is being traced

jit
1

I would like to call a different function when tracing. I have registered it through
torch::jit::RegisterOperators() and had to change it a little for that, so would like to call it only if it is tracing.
Is there a way to find if the module is being traced?
I currently rely on the fact that tenssor.shape[0] is a tensor while tracing, but an int otherwise. That is not the best way.

Something like torch.jit.is_running

1 Like
2

Hi thanks for posting the question. We do have a method do know if it’s in tracing state or not using torch._C._get_tracing_state(), it will return true if you are in tracing otherwise false. But note that this is a internal API and it might break in the future (also unlikely).

1 Like
3

Thanks! that would be very useful, a public facing API would be most very welcome. It will be better than relying on if isinstance(bbs.shape[0], torch.Tensor).