Avoid autocast deprecation warning in DataParallel (#130660)

awaelchli · guangyey · mlazos · commit 3f524bddc08c · 2024-07-17T19:09:50.000-07:00
Fixes #130659 Co-authored-by: Yu, Guangye <106960996+guangyey@users.noreply.github.com> Pull Request resolved: #130660 Approved by: https://github.com/guangyey, https://github.com/fegin, https://github.com/albanD
diff --git a/torch/nn/parallel/parallel_apply.py b/torch/nn/parallel/parallel_apply.py
@@ -4,7 +4,6 @@
 import torch
 from torch._utils import ExceptionWrapper
 from torch.cuda._utils import _get_device_index
-from torch.cuda.amp import autocast
 
 from ..modules import Module
 
@@ -89,9 +88,9 @@ def _worker(
         if stream is None:
             stream = torch.cuda.current_stream(device)
         try:
-            with torch.cuda.device(device), torch.cuda.stream(stream), autocast(
-                enabled=autocast_enabled
-            ):
+            with torch.cuda.device(device), torch.cuda.stream(
+                stream
+            ), torch.amp.autocast("cuda", enabled=autocast_enabled):
                 # this also avoids accidental slicing of `input` if it is a Tensor
                 if not isinstance(input, (list, tuple)):
                     input = (input,)
