Since you are assuming your custom CUDA extension is causing the trouble, I would probably try to come up with a minimal standalone code snippet using your module only (forward and backward) and check the allocations via print(torch.cuda.memory_summary()) in each iteration. If you are seeing a larger memory usage, then it would indeed point to your custom code. In this case, check first if you are storing any tensors, which would disallow PyTorch to free them. Also, I assume you are using libtorch or are you manually trying to manage the CUDA memory via raw malloc/free calls?