guard on whisper

huggingface · Oct 5, 2023 · 21f3b07 · 21f3b07
1 parent 9ac3402
commit 21f3b07
Showing 1 changed file with 7 additions and 5 deletions.
diff --git a/optimum/exporters/onnx/__main__.py b/optimum/exporters/onnx/__main__.py
@@ -348,11 +348,13 @@ def main_export(
         device=device,
         library_name=library_name,
     )
-    import torch
-    model = model.eval()
-    with torch.no_grad():
-        for i in range(len(model.model.decoder.layers)):
-            model.model.decoder.layers[i].encoder_attn = torch.jit.script(model.model.decoder.layers[i].encoder_attn)
+
+    if model.config.model_type == "whisper":
+        import torch
+        model = model.eval()
+        with torch.no_grad():
+            for i in range(len(model.model.decoder.layers)):
+                model.model.decoder.layers[i].encoder_attn = torch.jit.script(model.model.decoder.layers[i].encoder_attn)
 
     custom_architecture = False
     is_stable_diffusion = "stable-diffusion" in task