I'm trying to implement a pytorch custom layer grid_sampler on GPU by Core ML Custom Operators. Both of its inputs, input and grid can be 5-D. My implementation of encodeToCommandBuffer, which is MLCustomLayer protocol's function, is shown below. According to my current attempts, both value of id<MTLTexture> input and id<MTLTexture> grid don't meet expectations. So i wonder can MTLTexture be used to store 5-D input tensor as inputs of encodeToCommandBuffer? Or can anybody help to show me how to use MTLTexture correctly here? Thanks a lot!

- (BOOL)encodeToCommandBuffer:(id<MTLCommandBuffer>)commandBuffer
            inputs:(NSArray<id<MTLTexture>> *)inputs
           outputs:(NSArray<id<MTLTexture>> *)outputs
            error:(NSError * _Nullable *)error {
  NSLog(@"Dispatching to GPU");
  NSLog(@"inputs count %lu", (unsigned long)inputs.count);
  NSLog(@"outputs count %lu", (unsigned long)outputs.count);
  id<MTLComputeCommandEncoder> encoder = [commandBuffer
      computeCommandEncoderWithDispatchType:MTLDispatchTypeSerial];
    assert(encoder != nil);
   
  id<MTLTexture> input = inputs[0];
  id<MTLTexture> grid = inputs[1];
  id<MTLTexture> output = outputs[0];

  NSLog(@"inputs shape %lu, %lu, %lu, %lu", (unsigned long)input.width, (unsigned long)input.height, (unsigned long)input.depth, (unsigned long)input.arrayLength);
  NSLog(@"grid shape %lu, %lu, %lu, %lu", (unsigned long)grid.width, (unsigned long)grid.height, (unsigned long)grid.depth, (unsigned long)grid.arrayLength);
  if (encoder)
  {
    [encoder setTexture:input atIndex:0];
    [encoder setTexture:grid atIndex:1];
    [encoder setTexture:output atIndex:2];
     
    NSUInteger wd = grid_sample_Pipeline.threadExecutionWidth;
    NSUInteger ht = grid_sample_Pipeline.maxTotalThreadsPerThreadgroup / wd;
    MTLSize threadsPerThreadgroup = MTLSizeMake(wd, ht, 1);
    MTLSize threadgroupsPerGrid = MTLSizeMake((input.width + wd - 1) / wd, (input.height + ht - 1) / ht, input.arrayLength);

    [encoder setComputePipelineState:grid_sample_Pipeline];
    [encoder dispatchThreadgroups:threadgroupsPerGrid threadsPerThreadgroup:threadsPerThreadgroup];
    [encoder endEncoding];
     
  }
  else
    return NO;

  *error = nil;
  return YES;
}