Does NBDTs support jit?

[Optimox]

Hi,

Great project and article, congratulations.

I'm wondering if, once the training is finished, the final NBDT can be fully exported as a JIT file?

It seems that everything is written in pytorch and inference can be done in a single forward pass so it should be fine but I'm still a bit worried about the tree part and the self.rules.forward_with_decisions(x). Have you ever tried to export one of your models as JIT? Does it work as is?

Thanks in advance for your answer.