I've added some more suggestions and comments. Happy to assist with replicating the requested changes to all areas, let me know!

Thanks so much for the suggestions! code updated. And for further replication, I will submit a follow-up PR today. And for the rest, I would be very happy if you could help.

Do you really need "BACKEND_RESET_MAX_MEMORY_ALLOCATED"? This function will all reset_peak_memory_stats anyway as can be seen from here: https://pytorch.org/docs/stable/generated/torch.cuda.reset_max_memory_allocated.html

Do you really need "BACKEND_RESET_MAX_MEMORY_ALLOCATED"? This function will all reset_peak_memory_stats anyway as can be seen from here: https://pytorch.org/docs/stable/generated/torch.cuda.reset_max_memory_allocated.html

We can look at this in another PR.

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

This looks very promising! Thank you!

Have you run the tests for which the changes are being introduced in this PR?

yes, I run the tests on xpu and except backend_reset_max_memory_allocated, all tests work. But I can fix it in another PR as @hlky mentioned earlier.

@faaany backend_reset_max_memory_allocated should be a no-op for non-cuda, if it's giving an error we'll need to fix that. I was referring to using reset_max_memory_allocated when that function has changed to also call reset_peak_memory_stats, we'll review that change in behavior later.

sure, let me update.

Hi @sayakpaul @hlky @DN6 , could you do final-round review? Since I will have follow-up PRs, it would be great that it could be merged this week. Thanks a lot!

Hi folks, any feedback?