neural-backed-decision-trees icon indicating copy to clipboard operation
neural-backed-decision-trees copied to clipboard

Published 20 hours ago • verified publisher icon alvinwan

Does NBDTs support jit?

Open Optimox opened this issue 3 years ago • 0 comments

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.

Optimox avatar Jan 03 '22 14:01 Optimox