TODO Add support for selecting device (gpu support) Support other input and output types in JITModules (IValue) such as Dicts and Lists Implement eval & train for JITModule Implement bindings for (at least some) optimizers