Fix inference code

This should now work with any PyTorch model (Optimus is the example
given in the source code), as well as any HuggingFace model (adjusted
the code to be independent of any model source).