bug

Author: xadupre

CHANGELOGS.rst

@@ -4,6 +4,8 @@ Change Logs
 0.8.10
 ++++++
 
+* :pr:`384`: add ``weights_only=False`` when using :func:`torch.load`
+
 0.8.9
 +++++
 

_doc/examples/plot_export_tiny_llm_method_generate.py

@@ -145,6 +145,7 @@ def generate_text(
 # The following lines are a condensed copy with less comments.
 
 # from HuggingFace
+print("----------------")
 MODEL_NAME = "arnir0/Tiny-LLM"
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
@@ -160,14 +161,15 @@ def generate_text(
     convert_after_n_calls=3,
     dynamic_batch_for={"input_ids", "attention_mask", "past_key_values"},
 )
+model.forward = lambda *args, **kwargs: forward_replacement(*args, **kwargs)
 
 # from HuggingFace again
-prompt = "Continue: it rains..."
+prompt = "Continue: it rains, what should I do?"
 inputs = tokenizer(prompt, return_tensors="pt")
 outputs = model.generate(
     input_ids=inputs["input_ids"],
     attention_mask=inputs["attention_mask"],
-    max_length=50,
+    max_length=100,
     temperature=1,
     top_k=50,
     top_p=0.95,

onnx_diagnostic/export/api.py

@@ -738,7 +738,9 @@ def check_discrepancies(
         :param verbose: verbosity
         :return: results, a list of dictionaries, ready to be consumed by a dataframe
         """
-        assert self._export_done, "The onnx export was not done."
+        assert (
+            self._export_done
+        ), f"The onnx export was not done, only {len(self._inputs)} were stored."
         assert os.path.exists(self._input_file), f"input file {self._input_file!r} not found"
         assert os.path.exists(
             self._output_file
@@ -768,11 +770,11 @@ def check_discrepancies(
             print(f"[method_to_onnx.check_discrepancies] register classes {classes}")
             print(f"[method_to_onnx.check_discrepancies] load {self._input_file!r}")
         with torch.serialization.safe_globals(classes):
-            inputs = torch.load(self._input_file)
+            inputs = torch.load(self._input_file, weights_only=False)
         if verbose:
             print(f"[method_to_onnx.check_discrepancies] load {self._output_file!r}")
         with torch.serialization.safe_globals(classes):
-            outputs = torch.load(self._output_file)
+            outputs = torch.load(self._output_file, weights_only=False)
         assert len(inputs) == len(outputs), (
             f"Unexpected number of inputs {len(inputs)} and outputs {len(outputs)}, "
             f"inputs={string_type(inputs, with_shape=True)}, "