* Various fixes to support torch script. * Add tests to ensure that the model is torch scriptable. * Update tests.