Hi @frameworksengineer, thank you for your reply! That's great to know. I applied the @autoreleasepool block as you suggested and this resolves the memory leak. Thanks again!

After a lot of experimentation and further debugging, I've found the issue; the input tensors to the custom layer must be at most a rank-4 tensor. For my case, one of my inputs was rank-5. This seemed to force the fallback function evaluateOnCPUWithInputs to be used. I have not found anywhere in the documentation mentioning this restriction (although I might have missed it). It could be very helpful for others in the future if Apple/CoreML included this in their documentation, or even an example detailing the restrictions.