Hey, Did you end up finishing this thing? I am interested in this functionality, specifically for pytorch LongTensors. Thanks!