From d8caf8c8e91f720bc103b1252664c01449da4f1e Mon Sep 17 00:00:00 2001 From: Ella Charlaix Date: Mon, 9 Dec 2024 17:54:17 +0100 Subject: [PATCH] fix custom module test --- tests/onnx/test_onnx_export_custom_module.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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