Not an expert but maybe these instructions will help you get started https://github.com/pytorch/tutorials/pull/1821
@d4l3k