Solved by jerryzh168 in post #3 
              
              
                we do have some non-public API to do this: https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/native_functions.yaml#L3862 and https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/native/native_functions.yaml#L3868 
but they we might change the API when we officially release …