Only need onnx export function in export

LibrePhotos · Nov 5, 2023 · 509edd2 · 509edd2
1 parent ddc1a67
commit 509edd2
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/api/im2txt/sample.py b/api/im2txt/sample.py
@@ -6,7 +6,6 @@
 from PIL import Image
 from torchvision import transforms
 
-from torch.onnx import export, dynamo_export
 from api.im2txt.model import DecoderRNN, EncoderCNN
 import onnxruntime as ort
 
@@ -163,6 +162,8 @@ def generate_caption(
         return sentence
 
     def export_onnx(self, encoder_output_path, decoder_output_path):
+        from torch.onnx import export, dynamo_export
+
         self.load_models()
 
         # Define a sample input image tensor (adjust input size as needed)