Fix (learned_round): disable return QuantTensor during float inference

[pablomlago]

Reason for this PR

ImageNet entry-point was not working when enabling "--learned-round", as the quantization of the model was not properly enabled/disabled, due to the omission of the call to disable_return_quant_tensor/restore_return_quant_tensor.

Changes Made in this PR

Incorporate missing calls to update the return_quant_tensor for the model submodels, as well as add unit test to replicate the issue.

Testing Summary

Create new unit test for the class DisableEnableQuantization.

Risk Highlight

• This PR includes code from another work (please detail).
• This PR contains API-breaking changes.
• This PR depends on work in another PR (please provide links/details).
• This PR introduces new dependencies (please detail).
• There are coverage gaps not covered by tests.
• Documentation updates required in subsequent PR.

Checklist

• Code comments added to any hard-to-understand areas, if applicable.
• Changes generate no new warnings.
• Updated any relevant tests, if applicable.
• No conflicts with destination dev branch.
• I reviewed my own code changes.
• Initial CI/CD passing.
• 1+ reviews given, and any review issues addressed and approved.
• Post-review full CI/CD passing.

[Giuseppe5]

PRs should always point to dev :)