Failed to save quantized model

[lockeregg]

I tried to run this example https://github.com/intel/neural-compressor/blob/master/examples/3.x_api/pytorch/cv/static_quant/main.py, and I got an error in https://github.com/intel/neural-compressor/blob/09d4f2d6fb1a6aa91874a0b87a967067800462cb/examples/3.x_api/pytorch/cv/static_quant/main.py#L220

error message:

  q_model.save(example_inputs=example_inputs, output_dir="./saved_results")
  File "C:\Users\JackTC_Li\AppData\Roaming\Python\Python38\site-packages\neural_compressor\torch\algorithms\pt2e_quant\save_load.py", line 38, in save
    quantized_ep = torch.export.export(model, example_inputs, dynamic_shapes=dynamic_shapes)
  File "D:\anaconda3\envs\nc_test\lib\site-packages\torch\export\__init__.py", line 174, in export
    return _export(
  File "D:\anaconda3\envs\nc_test\lib\site-packages\torch\export\_trace.py", line 945, in wrapper
    raise e
  File "D:\anaconda3\envs\nc_test\lib\site-packages\torch\export\_trace.py", line 928, in wrapper
    ep = fn(*args, **kwargs)
  File "D:\anaconda3\envs\nc_test\lib\site-packages\torch\export\exported_program.py", line 89, in wrapper
    return fn(*args, **kwargs)
  File "D:\anaconda3\envs\nc_test\lib\site-packages\torch\export\_trace.py", line 1455, in _export
    aten_export_artifact = export_func(
  File "D:\anaconda3\envs\nc_test\lib\site-packages\torch\export\_trace.py", line 1060, in _strict_export
    gm_torch_level = _export_to_torch_ir(
  File "D:\anaconda3\envs\nc_test\lib\site-packages\torch\export\_trace.py", line 512, in _export_to_torch_ir
    gm_torch_level, _ = torch._dynamo.export(
  File "D:\anaconda3\envs\nc_test\lib\site-packages\torch\_dynamo\eval_frame.py", line 1379, in inner
    result_traced = opt_f(*args, **kwargs)
  File "D:\anaconda3\envs\nc_test\lib\site-packages\torch\nn\modules\module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "D:\anaconda3\envs\nc_test\lib\site-packages\torch\nn\modules\module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "D:\anaconda3\envs\nc_test\lib\site-packages\torch\_dynamo\eval_frame.py", line 433, in _fn
    return fn(*args, **kwargs)
  File "D:\anaconda3\envs\nc_test\lib\site-packages\torch\fx\graph_module.py", line 738, in call_wrapped
    return self._wrapped_call(self, *args, **kwargs)
  File "D:\anaconda3\envs\nc_test\lib\site-packages\torch\fx\graph_module.py", line 316, in __call__
    raise e
  File "D:\anaconda3\envs\nc_test\lib\site-packages\torch\fx\graph_module.py", line 303, in __call__
    return super(self.cls, obj).__call__(*args, **kwargs)  # type: ignore[misc]
  File "D:\anaconda3\envs\nc_test\lib\site-packages\torch\nn\modules\module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "D:\anaconda3\envs\nc_test\lib\site-packages\torch\nn\modules\module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
  File "D:\anaconda3\envs\nc_test\lib\site-packages\torch\_dynamo\convert_frame.py", line 1116, in __call__
    return self._torchdynamo_orig_callable(
  File "D:\anaconda3\envs\nc_test\lib\site-packages\torch\_dynamo\convert_frame.py", line 472, in __call__
    return _compile(
  File "D:\anaconda3\envs\nc_test\lib\site-packages\torch\_utils_internal.py", line 84, in wrapper_function
    return StrobelightCompileTimeProfiler.profile_compile_time(
  File "D:\anaconda3\envs\nc_test\lib\site-packages\torch\_strobelight\compile_time_profiler.py", line 129, in profile_compile_time
    return func(*args, **kwargs)
  File "D:\anaconda3\envs\nc_test\lib\contextlib.py", line 75, in inner
    return func(*args, **kwds)
  File "D:\anaconda3\envs\nc_test\lib\site-packages\torch\_dynamo\convert_frame.py", line 817, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
  File "D:\anaconda3\envs\nc_test\lib\site-packages\torch\_dynamo\utils.py", line 231, in time_wrapper
    r = func(*args, **kwargs)
  File "D:\anaconda3\envs\nc_test\lib\site-packages\torch\_dynamo\convert_frame.py", line 636, in compile_inner
    out_code = transform_code_object(code, transform)
  File "D:\anaconda3\envs\nc_test\lib\site-packages\torch\_dynamo\bytecode_transformation.py", line 1185, in transform_code_object
    transformations(instructions, code_options)
  File "D:\anaconda3\envs\nc_test\lib\site-packages\torch\_dynamo\convert_frame.py", line 178, in _fn
    return fn(*args, **kwargs)
  File "D:\anaconda3\envs\nc_test\lib\site-packages\torch\_dynamo\convert_frame.py", line 582, in transform
    tracer.run()
  File "D:\anaconda3\envs\nc_test\lib\site-packages\torch\_dynamo\symbolic_convert.py", line 2451, in run
    super().run()
  File "D:\anaconda3\envs\nc_test\lib\site-packages\torch\_dynamo\symbolic_convert.py", line 893, in run
    while self.step():
  File "D:\anaconda3\envs\nc_test\lib\site-packages\torch\_dynamo\symbolic_convert.py", line 805, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "D:\anaconda3\envs\nc_test\lib\site-packages\torch\_dynamo\symbolic_convert.py", line 499, in wrapper
    return inner_fn(self, inst)
  File "D:\anaconda3\envs\nc_test\lib\site-packages\torch\_dynamo\symbolic_convert.py", line 1459, in CALL_FUNCTION
    self.call_function(fn, args, {})
  File "D:\anaconda3\envs\nc_test\lib\site-packages\torch\_dynamo\symbolic_convert.py", line 743, in call_function
    self.push(fn.call_function(self, args, kwargs))
  File "D:\anaconda3\envs\nc_test\lib\site-packages\torch\_dynamo\variables\functions.py", line 665, in call_function
    unimplemented(msg)
  File "D:\anaconda3\envs\nc_test\lib\site-packages\torch\_dynamo\exc.py", line 221, in unimplemented
    raise Unsupported(msg)
torch._dynamo.exc.Unsupported: 'skip function tree_flatten_spec in file D:\anaconda3\envs\nc_test\lib\site-packages\torch\fx\_pytree.py'

from user code:
   File "<eval_with_key>.14", line 5, in forward
    arg0, = fx_pytree.tree_flatten_spec(([x], {}), self._in_spec)

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information

system settings: OS: Windows 11 Pytorch: 2.2~2.4.0 Neural compressor: 3.0, 3.1, 3.1.dev14+g09d4f2d6fb

Thank you for your help.