reduce some waring for onnx export, add test code

Tlntin · Jul 29, 2024 · 7ae6095 · 7ae6095
1 parent 9eeea23
commit 7ae6095
Show file tree

Hide file tree

Showing 5 changed files with 390 additions and 133 deletions.
diff --git a/export/export_onnx.py b/export/export_onnx.py
@@ -150,12 +150,12 @@ def export_onnx(
         attention_mask,
         position_ids,
         past_key_values,
-        None,  # inputs_embeds: Optional[torch.FloatTensor] = None,
-        None,  # labels: Optional[torch.LongTensor] = None,
-        True,  # use_cache: Optional[bool] = None,
+        # None,  # inputs_embeds: Optional[torch.FloatTensor] = None,
+        # None,  # labels: Optional[torch.LongTensor] = None,
+        # True,  # use_cache: Optional[bool] = None,
         True,  # output_attentions: Optional[bool] = None,
-        None,  # output_hidden_states
-        False  # return_dict:
+        # None,  # output_hidden_states
+        # False  # return_dict:
     )
     model.eval()
     with torch.no_grad():