[Pytorch Mobile] Preserved all functions generated by bundled inputs (#51496)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/51496
A previous change added the possibility of more functions being generated when bundled inputs are attached. Want to preserve those here in optimize_for_mobile
ghstack-source-id: 120862718
Test Plan:
Created a dummy model. Augment several methods with bundled inputs. Call optimize for mobile. Verified the functions are still there.
Discovered a weird interaction between freeze_module and bundled inputs. If the user does something like
inputs =[<inputs>]
augment_many_model_functions_with_bundled_inputs(
model,
inputs={
model.forward : inputs,
model.foo : inputs,
}
)
to attach their bundled inputs, freeze_module within optimize_for_mobile will error out. Instead the user would need to do something like
inputs =[<inputs>]
inputs2 =[<inputs>] # Nominally the same as the the inputs above
augment_many_model_functions_with_bundled_inputs(
model,
inputs={
model.forward : inputs,
model.foo : inputs2,
}
)
Reviewed By: dhruvbird
Differential Revision: D26005708
fbshipit-source-id: 3e908c0f7092a57da9039fbc395aee6bf9dd2b20
