Add lennard-jones model (#1180)
Summary:
The main thing we want to test here is that using functorch's vmap x jacrev on some fully connected layers works. This is a common sequence of operations for functorch and this PR is an example of one such model that exercises it.
Pull Request resolved: https://github.com/pytorch/benchmark/pull/1180
Test Plan:
- `python test.py -k "test_lennard_jones_eval_cuda"`
- `python test.py -k "test_lennard_jones_train_cuda"`
Reviewed By: xuzhao9
Differential Revision: D39670769
Pulled By: zou3519
fbshipit-source-id: e40d807cbd5c8982b43c92e33696d81fc08101a9
