See related issue on tensorflow github:
Custom shuffle layer leaks memory when run on Apple M1 GPU with tensorflow-metal
Add a Comment
See related issue on tensorflow github:
Hi @sirno
Thanks for posting the issue and the provided sample code. Based on it I was able to track the leaking memory to the implementation of the TopK operation which ends up being called underneath the argsort op and added the necessary fixes to prevent the leaks. This fix will come out in the next release of the tensorflow-metal plugin.