Add a new "test" argument to models (#735)
Summary:
This PR adds a new argument, 'test', which can be 'train' and 'eval', in the model initialization function. It has the following benefits:
- At model initialization, we can choose to only initialize "train" or "eval" model datasets, to avoid wasting device memory
- Separate the space of arguments. It is common that train and eval tests have different experimental features to apply. For example, optimize_for_inference, and tensorrt only work for eval.
- Add an "extra_args" argument to the model, to enable more optional features in the future.
Pull Request resolved: https://github.com/pytorch/benchmark/pull/735
Reviewed By: erichan1
Differential Revision: D34033140
Pulled By: xuzhao9
fbshipit-source-id: 449a963b1ee7ef450b0a5ccb24718d394dd8be40