Enable TE in test_jit.py (#44200)
Summary:
Enable TE in test_jit.py and adjust/fix tests accordingly.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/44200
Reviewed By: SplitInfinity
Differential Revision: D23673624
Pulled By: Krovatkin
fbshipit-source-id: 5999725c7aacc6ee77885eb855a41ddfb4d9a8d8
