Add torch.jit.trace option

[tugsbayasgalan]

Stack from ghstack (oldest at bottom):

• -> #2255

python run.py resnet50 -d cuda -t train --backend torchscript_trace

Differential Revision: D56849015

[tugsbayasgalan]

@tugsbayasgalan has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.