Fix AMP autocast deprecation warning for PyTorch>=2.5

detectron2/engine/train_loop.py

@@ -479,7 +479,7 @@ def run_step(self):
         """
         assert self.model.training, "[AMPTrainer] model was changed to eval mode!"
         assert torch.cuda.is_available(), "[AMPTrainer] CUDA is required for AMP training!"
-        from torch.cuda.amp import autocast
+        from detectron2.utils.torch_amp import autocast
 
         start = time.perf_counter()
         data = next(self._data_loader_iter)

detectron2/utils/torch_amp.py

@@ -0,0 +1,8 @@
+try:
+    from torch.amp import autocast as _autocast
+
+    def autocast(dtype=None):
+        return _autocast("cuda", dtype=dtype)
+
+except ImportError:
+    from torch.cuda.amp import autocast  # torch < 2.0