Fix the deprecated torch.cuda.amp module (#21)

* Fix the deprecated torch.cuda.amp module

* Fix test according to the latest GradScaler definition.

Author: MGAMZ

mmengine/optim/optimizer/amp_optimizer_wrapper.py

@@ -1,5 +1,6 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 from contextlib import contextmanager
+from functools import partial
 from typing import Union
 
 import torch
@@ -17,7 +18,8 @@
 elif is_mlu_available():
     from torch.mlu.amp import GradScaler
 else:
-    from torch.cuda.amp import GradScaler
+    from torch.amp import GradScaler
+    GradScaler = partial(GradScaler, device='cuda')
 
 
 @OPTIM_WRAPPERS.register_module()

tests/test_optim/test_optimizer/test_optimizer_wrapper.py

@@ -8,7 +8,7 @@
 import torch.distributed as torch_dist
 import torch.nn as nn
 from parameterized import parameterized
-from torch.cuda.amp import GradScaler
+from torch.amp import GradScaler
 from torch.nn.parallel.distributed import DistributedDataParallel
 from torch.optim import SGD, Adam, Optimizer