One thing to keep in mind is that you likely need to use a very new version of PyTorch (i.e. compiled from the master git branch or nightly releases).
The example you linked was added just a week ago, so it might need features from very recent PyTorch. The background is that some bits of the C++ API still have some catch-up to do to achieve parity with “classic” PyTorch.
Best regards
Thomas