♨

Author: xxaier

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',