Complete re-write for pytorch 0.2.0