♨

FlagAI-Open · Jun 4, 2023 · 57d3b07 · 57d3b07
1 parent 6044594
commit 57d3b07
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/onnx/export/onnx_export.py b/onnx/export/onnx_export.py
@@ -36,9 +36,12 @@ def onnx_export(model, args, **kwds):
 
 # 参考 https://github.com/OFA-Sys/Chinese-CLIP/blob/master/cn_clip/deploy/pytorch_to_onnx.py
 def export(txt, img):
+
   onnx_export(txt,
               tokenizer(['a photo of cat', 'a image of cat'], ),
               input_names=['input', 'attention_mask'],
+              input_types=[torch.uint32],
+              output_types=[torch.uint32],
               dynamic_axes={
                   'input': {
                       0: 'batch',