I'm trying to implement a 5-D input tensor version of custom grid_sample based on this work. For encoding one of the input tensor, grid to a MTLTexture as shown in officical guide, i need to transpose grid from (N×D_grid×W_grid×H_grid×3) to (N×3×D_grid×W_grid ×H_grid) by builder.add_transpose similar like shown here. But in my implementation, i find that adding this transpose op will make the custom layer always run on CPU. Without this layer, the data can get in GPU. I doubt will builder.add_transpose have an effect on this?
System Information
- xcode: 13.2
- coremltools: 4.1/5.1.0
- test device: iphone 11