Solution to make allennlp jit compilable

[EmanueleGusso]

Hi everybody,
I tried to jit compile the model wrapped in CorefPredictor (https://github.com/allenai/allennlp-models/blob/main/allennlp_models/coref/models/coref.py) but I get some warnings. In fact, as it is, the model cannot be jit compiled because in the forward method there are if-then-else statement and sometimes the torch tensors are transformed into numpy arrays to perform some operations and then they are reconverted into torch tensors.
So I have two questions:

1. will this feature be implemented?
2. if I wanted to try to implement this feature, would it be feasible? I mean, for example, one solution could be to split the model in two when the numpy transformations or the control-flow happen. Does someone of you already know if this is possible or if this is not possible due to the particular flow that the forward of this model have (or for some other reason that I do not know)?

This would be very important for me. I hope you can help :)
I thank you in advance for your availability.

[dirkgr]

We are discussing this at #5565.