[Chore] Cleanup mem_utils.py

[DarkLight1337]

Purpose

• Ensure MemoryProfilingResult snapshots all have the same device
• Remove unnecessary noqas

Some of the changes are taken from #22070

Test Plan

Test Result

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

[gemini-code-assist]

Code Review

This pull request aims to clean up mem_utils.py by ensuring memory snapshots use a consistent device and removing unnecessary noqa comments. While the intention to enforce device consistency is good, the implementation is flawed. The MemorySnapshot class stores the actual device in a private device_ attribute, but the changes incorrectly attempt to use the public device attribute, which can be None. This could lead to incorrect device selection in a multi-device environment. I've provided two high-severity comments with suggestions to use the device_ attribute to fix this correctness issue. The other changes, such as refactoring and removing noqa comments, are good improvements.