diff --git a/tests/onnx/test_onnx_export_custom_module.py b/tests/onnx/test_onnx_export_custom_module.py index 4398c14f01d..9416093c841 100644 --- a/tests/onnx/test_onnx_export_custom_module.py +++ b/tests/onnx/test_onnx_export_custom_module.py @@ -22,7 +22,7 @@ if is_torch_available(): import torch - from transformers.models.deberta import modeling_deberta + from transformers.models.sew_d import modeling_sew_d from optimum.utils import check_if_torch_greater @@ -36,7 +36,7 @@ def test_training(self): """Tests export of StableDropout in training mode.""" devnull = open(os.devnull, "wb") # drop_prob must be > 0 for the test to be meaningful - sd = modeling_deberta.StableDropout(0.1) + sd = modeling_sew_d.StableDropout(0.1) # Avoid warnings in training mode do_constant_folding = False # Dropout is a no-op in inference mode