Use torch.amp.autocast instead of deprecated torch.cuda.amp.autocast …

…in nemotron Signed-off-by: Dmitry Rogozhkin <dmitry.v.rogozhkin@intel.com>
huggingface · Nov 26, 2024 · 86f8b24 · 86f8b24
1 parent cbcf4ff
commit 86f8b24
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/transformers/models/nemotron/modeling_nemotron.py b/src/transformers/models/nemotron/modeling_nemotron.py
@@ -76,7 +76,7 @@ def __init__(
 
     def forward(self, input: Tensor) -> Tensor:
         args = _cast_if_autocast_enabled(input, self.normalized_shape, self.weight + 1, self.bias, self.eps)
-        with torch.cuda.amp.autocast(enabled=False):
+        with torch.amp.autocast(input.device.type, enabled=False):
             return F.layer_norm(*args)