lazy import torch._softmax_backward_data for better compatibility (#20796)

lazy import torch._softmax_backward_data
Signed-off-by: daquexian <daquexian566@gmail.com>
Signed-off-by: daquexian <daquexian566@gmail.com>

src/transformers/pytorch_utils.py

@@ -16,7 +16,7 @@ from typing import Callable, List, Optional, Set, Tuple, Union
 
 import torch
 from packaging import version
-from torch import _softmax_backward_data, nn
+from torch import nn
 
 from .utils import logging
 
@@ -48,6 +48,8 @@ def softmax_backward_data(parent, grad_output, output, dim, self):
     to the torch version detected.
     """
 
+    from torch import _softmax_backward_data
+
     if is_torch_less_than_1_11:
         return _softmax_backward_data(grad_output, output, parent.dim, self)
     else: