fix custom module test

huggingface · Dec 9, 2024 · d8caf8c · d8caf8c
1 parent 74e0bad
commit d8caf8c
Showing 1 changed file with 2 additions and 2 deletions.
diff --git 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