Fix the deprecated torch.cuda.amp module

[MGAMZ]

This sub-PR is related to open-mmlab#1665

Brief

According to PyTorch:

torch.cuda.amp.GradScaler(args...) is deprecated. Please use torch.amp.GradScaler("cuda", args...) instead.

This includes two related replacement:

1. amp_optimizer_wrapper
2. test_optimizer_wrapper

PyTest Result After this PR

pytest tests/test_optim/test_optimizer/test_optimizer_wrapper.py

===================================================================== test session starts ======================================================================
platform linux -- Python 3.13.9, pytest-8.4.2, pluggy-1.6.0
rootdir: /home/mgam/mgam_repos/onedl-mmengine
configfile: pyproject.toml
plugins: hydra-core-1.3.2, anyio-4.11.0
collected 37 items                                                                                                                                             

tests/test_optim/test_optimizer/test_optimizer_wrapper.py sssssssssssssssssss..................                                                          [100%]

======================================================================= warnings summary =======================================================================
../../miniforge3/envs/onedl/lib/python3.13/site-packages/torch/cuda/__init__.py:63
  /home/mgam/miniforge3/envs/onedl/lib/python3.13/site-packages/torch/cuda/__init__.py:63: FutureWarning: The pynvml package is deprecated. Please install nvidia-ml-py instead. If you did not install pynvml directly, please report this to the maintainers of the package that installed pynvml for you.
    import pynvml  # type: ignore[import]

mmengine/utils/misc.py:477
  /home/mgam/mgam_repos/onedl-mmengine/mmengine/utils/misc.py:477: DeprecationWarning: 'maxsplit' is passed as positional argument
    summary_and_body = re.split(pattern, docstring, 1)

tests/test_optim/test_optimizer/test_optimizer_wrapper.py::TestAmpOptimWrapper::test_optim_context_4
  /home/mgam/mgam_repos/onedl-mmengine/mmengine/runner/amp.py:119: DeprecationWarning: torch.get_autocast_gpu_dtype() is deprecated. Please use torch.get_autocast_dtype('cuda') instead. (Triggered internally at /pytorch/torch/csrc/autograd/init.cpp:856.)
    dtype = torch.get_autocast_gpu_dtype()

-- Docs: https://docs.pytest.org/en/stable/how-to/capture-warnings.html
==================================================== 18 passed, 19 skipped, 3 warnings in 67.40s (0:01:07) =====================================================

[Copilot]

Pull Request Overview

This PR modernizes the GradScaler import by migrating from the deprecated torch.cuda.amp.GradScaler to the newer torch.amp.GradScaler API, which is device-agnostic. The changes ensure backward compatibility by wrapping the new API with a partial function that defaults to CUDA device when neither NPU nor MLU is available.

• Updated GradScaler import from torch.cuda.amp to torch.amp in production code and tests
• Added a device-specific wrapper using functools.partial to maintain CUDA as default device

Reviewed Changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
mmengine/optim/optimizer/amp_optimizer_wrapper.py	Updated GradScaler import to use torch.amp and added partial wrapper for CUDA device specification
tests/test_optim/test_optimizer/test_optimizer_wrapper.py	Updated test imports to use the new torch.amp.GradScaler API

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[read-the-docs-community]

Documentation build overview

📚 onedl-mmengine | 🛠️ Build #30193815 | 📁 Comparing 5e31726 against latest (e529128)

🔍 Preview build

Show files changed (2 files in total): 📝 2 modified | ➕ 0 added | ➖ 0 deleted

File	Status
_modules/mmengine/_strategy/base.html	📝 modified
_modules/mmengine/optim/optimizer/amp_optimizer_wrapper.html	📝 modified

[lauriebax]

If you fix the linting error I will merge.

[MGAMZ]

@lauriebax Lint fixed. And, as described in open-mmlab#1676:

yapf=0.32.0 requires lib2to3 which is deprecated since python=3.11, and is removed in python=3.13. 2to3 Doc

This issue cause contributors being unable to run yapf lint in python>=3.13. So I suggest upgrade the yapf version in pre-commit which no longer requires this deprecated lib.

I suggest to update yapf version in pre-commit. If ONEDL-mmengine accept this idea, I can help to create a PR.