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

…in nemotron Signed-off-by: Dmitry Rogozhkin <[email protected]>
dvrogozh · Nov 27, 2024 · 0191f49 · 0191f49
1 parent d476d8d
commit 0191f49
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)