What do you mean by if I use self.training? My problem is that inference_training might contain raw python code that is not exportable. But introducing a branch makes the script_method not exportable
The question is:
Let’s say I have module A and module B. module A uses module B.
module A code is not branchy (inference and training is the same). I want to be able to call module A with just forward. I don’t want module A to have to have a forward_training vs forward_inference.
module B code is branchy (inference / training are different), I am ok for it to have a forward_training and a forward_inference.