For certain models, the forward path might be slightly different in training vs inference.
If I need to export the inference code (but not the training code). What is a good design pattern to achieve this.
@jit.script_method
forward(self, inputs):
   if TRAINING:
       /// training code that does not need to be compiled
   else:
      /// inference code, compile this