pytorch - ✅(Solved) Fix `torch.export.save` fails when `torch.uint32` is present [1 pull requests, 1 comments, 1 participants]

pytorch2026-04-05 22:49:45

ON THIS PAGE

Recommended Tools

×6

Utilities matched from this issue’s tags and category — try them while you read without losing context.

GitHub issue graph ai analysis

Paste a GitHub issue URL. We fetch that issue, discover linked issues from bodies/comments/timeline, collect linked pull requests, and produce a structured English report.

The report is written in English Markdown for sharing and archival.

GitHub issue URL

Helpful · Quick feedback

GitHub stats

pytorch/pytorch#179433•Fetched 2026-04-08 02:51:43

View on GitHub

Comments

Participants

Timeline

Reactions

Author

ar0ck

Participants

ar0ck

Timeline (top)

mentioned ×9subscribed ×9labeled ×4commented ×1

Error Message

Traceback (most recent call last): File "/.../lib/python3.14/site-packages/torch/export/serde/serialize.py", line 2125, in serialize_graph getattr(self, f"handle{node.op}")(node) ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^ File "/.../lib/python3.14/site-packages/torch/_export/serde/serialize.py", line 814, in handle_call_function inputs=self.serialize_inputs(node.target, node.args, node.kwargs), ~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/.../lib/python3.14/site-packages/torch/_export/serde/serialize.py", line 1181, in serialize_inputs arg=self.serialize_input( ~~~~~~~~~~~~~~~~~~~~^ kwargs[schema_arg.name], schema_arg.type ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ ), ^ File "/.../lib/python3.14/site-packages/torch/_export/serde/serialize.py", line 1534, in serialize_input return Argument.create(as_scalar_type=_TORCH_TO_SERIALIZE_DTYPE[arg]) ~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^ KeyError: torch.uint32

The above exception was the direct cause of the following exception:

Traceback (most recent call last): File "bug.py", line 8, in <module> torch.export.save(ep, "ep.pt2") ~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^ File "/.../lib/python3.14/site-packages/torch/export/init.py", line 274, in save package_pt2( ~~~~~~~~~~~^ f, ^^ ...<3 lines>... opset_version=opset_version, ^^^^^^^^^^^^^^^^^^^^^^^^^^^^ ) ^ File "/.../lib/python3.14/site-packages/torch/export/pt2_archive/_package.py", line 703, in package_pt2 _package_exported_programs( ~~~~~~~~~~~~~~~~~~~~~~~~~~^ archive_writer, exported_programs, pickle_protocol=pickle_protocol ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ ) ^ File "/.../lib/python3.14/site-packages/torch/export/pt2_archive/_package.py", line 602, in _package_exported_programs artifact: SerializedArtifact = serialize( ~~~~~~~~~^ ep, ^^^ opset_version, ^^^^^^^^^^^^^^ pickle_protocol, ^^^^^^^^^^^^^^^^ ) ^ File "/.../lib/python3.14/site-packages/torch/_export/serde/serialize.py", line 3605, in serialize ).serialize(exported_program) ~~~~~~~~~^^^^^^^^^^^^^^^^^^ File "/.../lib/python3.14/site-packages/torch/_export/serde/serialize.py", line 2194, in serialize serialized_graph_module = gm_serializer.serialize(exported_program.graph_module) File "/.../lib/python3.14/site-packages/torch/_export/serde/serialize.py", line 2158, in serialize graph = self.serialize_graph(graph_module) File "/.../lib/python3.14/site-packages/torch/_export/serde/serialize.py", line 2127, in serialize_graph raise SerializeError( f"Failed serializing node {node} in graph: {node.format_node()}\n Original exception {traceback.format_exc()}" ) from e torch._export.serde.serialize.SerializeError: Failed serializing node empty in graph: %empty : [num_users=1] = call_function[target=torch.ops.aten.empty.memory_format](args = ([],), kwargs = {dtype: torch.uint32, device: cpu, pin_memory: False}) Original exception Traceback (most recent call last): File "/.../lib/python3.14/site-packages/torch/export/serde/serialize.py", line 2125, in serialize_graph getattr(self, f"handle{node.op}")(node) ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^ File "/.../lib/python3.14/site-packages/torch/_export/serde/serialize.py", line 814, in handle_call_function inputs=self.serialize_inputs(node.target, node.args, node.kwargs), ~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/.../lib/python3.14/site-packages/torch/_export/serde/serialize.py", line 1181, in serialize_inputs arg=self.serialize_input( ~~~~~~~~~~~~~~~~~~~~^ kwargs[schema_arg.name], schema_arg.type ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ ), ^ File "/.../lib/python3.14/site-packages/torch/_export/serde/serialize.py", line 1534, in serialize_input return Argument.create(as_scalar_type=_TORCH_TO_SERIALIZE_DTYPE[arg]) ~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^ KeyError: torch.uint32

Fix Action

Fixed

Fixed by PR: Support torch.uint{32,64} in torch.export.save (https://github.com/pytorch/pytorch/pull/179434)

PR fix notes

PR #179434: Support `torch.uint{32,64}` in `torch.export.save`

Repository: pytorch/pytorch
Author: ar0ck
State: open | merged: False
Link: https://github.com/pytorch/pytorch/pull/179434

Description (problem / solution / changelog)

Closes #179433.

Changed files

test/export/test_serialize.py (modified, +3/-1)
torch/_export/serde/export_schema.thrift (modified, +3/-1)
torch/_export/serde/schema.py (modified, +5/-1)
torch/_export/serde/schema.yaml (modified, +4/-2)
torch/_export/serde/serialize.py (modified, +2/-0)
torch/csrc/inductor/aoti_torch/generated_enum_converters.h (modified, +2/-0)
torch/csrc/utils/generated_serialization_types.h (modified, +7/-1)

Code Example

import torch

class Module(torch.nn.Module):
    def forward(self):
        return torch.empty((), dtype=torch.uint32)

ep = torch.export.export(Module(), ())
torch.export.save(ep, "ep.pt2")

---

Traceback (most recent call last):
  File "/.../lib/python3.14/site-packages/torch/_export/serde/serialize.py", line 2125, in serialize_graph
    getattr(self, f"handle_{node.op}")(node)
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^
  File "/.../lib/python3.14/site-packages/torch/_export/serde/serialize.py", line 814, in handle_call_function
    inputs=self.serialize_inputs(node.target, node.args, node.kwargs),
           ~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/.../lib/python3.14/site-packages/torch/_export/serde/serialize.py", line 1181, in serialize_inputs
    arg=self.serialize_input(
        ~~~~~~~~~~~~~~~~~~~~^
        kwargs[schema_arg.name], schema_arg.type
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    ),
    ^
  File "/.../lib/python3.14/site-packages/torch/_export/serde/serialize.py", line 1534, in serialize_input
    return Argument.create(as_scalar_type=_TORCH_TO_SERIALIZE_DTYPE[arg])
                                          ~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^
KeyError: torch.uint32

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "bug.py", line 8, in <module>
    torch.export.save(ep, "ep.pt2")
    ~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^
  File "/.../lib/python3.14/site-packages/torch/export/__init__.py", line 274, in save
    package_pt2(
    ~~~~~~~~~~~^
        f,
        ^^
    ...<3 lines>...
        opset_version=opset_version,
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File "/.../lib/python3.14/site-packages/torch/export/pt2_archive/_package.py", line 703, in package_pt2
    _package_exported_programs(
    ~~~~~~~~~~~~~~~~~~~~~~~~~~^
        archive_writer, exported_programs, pickle_protocol=pickle_protocol
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File "/.../lib/python3.14/site-packages/torch/export/pt2_archive/_package.py", line 602, in _package_exported_programs
    artifact: SerializedArtifact = serialize(
                                   ~~~~~~~~~^
        ep,
        ^^^
        opset_version,
        ^^^^^^^^^^^^^^
        pickle_protocol,
        ^^^^^^^^^^^^^^^^
    )
    ^
  File "/.../lib/python3.14/site-packages/torch/_export/serde/serialize.py", line 3605, in serialize
    ).serialize(exported_program)
      ~~~~~~~~~^^^^^^^^^^^^^^^^^^
  File "/.../lib/python3.14/site-packages/torch/_export/serde/serialize.py", line 2194, in serialize
    serialized_graph_module = gm_serializer.serialize(exported_program.graph_module)
  File "/.../lib/python3.14/site-packages/torch/_export/serde/serialize.py", line 2158, in serialize
    graph = self.serialize_graph(graph_module)
  File "/.../lib/python3.14/site-packages/torch/_export/serde/serialize.py", line 2127, in serialize_graph
    raise SerializeError(
        f"Failed serializing node {node} in graph: {node.format_node()}\n Original exception {traceback.format_exc()}"
    ) from e
torch._export.serde.serialize.SerializeError: Failed serializing node empty in graph: %empty : [num_users=1] = call_function[target=torch.ops.aten.empty.memory_format](args = ([],), kwargs = {dtype: torch.uint32, device: cpu, pin_memory: False})
 Original exception Traceback (most recent call last):
  File "/.../lib/python3.14/site-packages/torch/_export/serde/serialize.py", line 2125, in serialize_graph
    getattr(self, f"handle_{node.op}")(node)
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^
  File "/.../lib/python3.14/site-packages/torch/_export/serde/serialize.py", line 814, in handle_call_function
    inputs=self.serialize_inputs(node.target, node.args, node.kwargs),
           ~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/.../lib/python3.14/site-packages/torch/_export/serde/serialize.py", line 1181, in serialize_inputs
    arg=self.serialize_input(
        ~~~~~~~~~~~~~~~~~~~~^
        kwargs[schema_arg.name], schema_arg.type
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    ),
    ^
  File "/.../lib/python3.14/site-packages/torch/_export/serde/serialize.py", line 1534, in serialize_input
    return Argument.create(as_scalar_type=_TORCH_TO_SERIALIZE_DTYPE[arg])
                                          ~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^
KeyError: torch.uint32

---

Collecting environment information...
PyTorch version: 2.12.0.dev20260402
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 26.3.1 (arm64)
GCC version: Could not collect
Clang version: 17.0.0 (clang-1700.6.3.2)
CMake version: version 4.2.3
Libc version: N/A

Python version: 3.14.2 (v3.14.2:df793163d58, Dec  5 2025, 12:18:06) [Clang 16.0.0 (clang-1600.0.26.6)] (64-bit runtime)
Python platform: macOS-26.3.1-arm64-arm-64bit-Mach-O
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
Is XPU available: False
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True
Caching allocator config: N/A

CPU:
Apple M4 Pro

Versions of relevant libraries:
[pip3] Could not collect
[conda] Could not collect

RAW_BUFFERClick to expand / collapse

🐛 Describe the bug

import torch

class Module(torch.nn.Module):
    def forward(self):
        return torch.empty((), dtype=torch.uint32)

ep = torch.export.export(Module(), ())
torch.export.save(ep, "ep.pt2")

Error logs

Traceback (most recent call last):
  File "/.../lib/python3.14/site-packages/torch/_export/serde/serialize.py", line 2125, in serialize_graph
    getattr(self, f"handle_{node.op}")(node)
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^
  File "/.../lib/python3.14/site-packages/torch/_export/serde/serialize.py", line 814, in handle_call_function
    inputs=self.serialize_inputs(node.target, node.args, node.kwargs),
           ~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/.../lib/python3.14/site-packages/torch/_export/serde/serialize.py", line 1181, in serialize_inputs
    arg=self.serialize_input(
        ~~~~~~~~~~~~~~~~~~~~^
        kwargs[schema_arg.name], schema_arg.type
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    ),
    ^
  File "/.../lib/python3.14/site-packages/torch/_export/serde/serialize.py", line 1534, in serialize_input
    return Argument.create(as_scalar_type=_TORCH_TO_SERIALIZE_DTYPE[arg])
                                          ~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^
KeyError: torch.uint32

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "bug.py", line 8, in <module>
    torch.export.save(ep, "ep.pt2")
    ~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^
  File "/.../lib/python3.14/site-packages/torch/export/__init__.py", line 274, in save
    package_pt2(
    ~~~~~~~~~~~^
        f,
        ^^
    ...<3 lines>...
        opset_version=opset_version,
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File "/.../lib/python3.14/site-packages/torch/export/pt2_archive/_package.py", line 703, in package_pt2
    _package_exported_programs(
    ~~~~~~~~~~~~~~~~~~~~~~~~~~^
        archive_writer, exported_programs, pickle_protocol=pickle_protocol
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File "/.../lib/python3.14/site-packages/torch/export/pt2_archive/_package.py", line 602, in _package_exported_programs
    artifact: SerializedArtifact = serialize(
                                   ~~~~~~~~~^
        ep,
        ^^^
        opset_version,
        ^^^^^^^^^^^^^^
        pickle_protocol,
        ^^^^^^^^^^^^^^^^
    )
    ^
  File "/.../lib/python3.14/site-packages/torch/_export/serde/serialize.py", line 3605, in serialize
    ).serialize(exported_program)
      ~~~~~~~~~^^^^^^^^^^^^^^^^^^
  File "/.../lib/python3.14/site-packages/torch/_export/serde/serialize.py", line 2194, in serialize
    serialized_graph_module = gm_serializer.serialize(exported_program.graph_module)
  File "/.../lib/python3.14/site-packages/torch/_export/serde/serialize.py", line 2158, in serialize
    graph = self.serialize_graph(graph_module)
  File "/.../lib/python3.14/site-packages/torch/_export/serde/serialize.py", line 2127, in serialize_graph
    raise SerializeError(
        f"Failed serializing node {node} in graph: {node.format_node()}\n Original exception {traceback.format_exc()}"
    ) from e
torch._export.serde.serialize.SerializeError: Failed serializing node empty in graph: %empty : [num_users=1] = call_function[target=torch.ops.aten.empty.memory_format](args = ([],), kwargs = {dtype: torch.uint32, device: cpu, pin_memory: False})
 Original exception Traceback (most recent call last):
  File "/.../lib/python3.14/site-packages/torch/_export/serde/serialize.py", line 2125, in serialize_graph
    getattr(self, f"handle_{node.op}")(node)
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^
  File "/.../lib/python3.14/site-packages/torch/_export/serde/serialize.py", line 814, in handle_call_function
    inputs=self.serialize_inputs(node.target, node.args, node.kwargs),
           ~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/.../lib/python3.14/site-packages/torch/_export/serde/serialize.py", line 1181, in serialize_inputs
    arg=self.serialize_input(
        ~~~~~~~~~~~~~~~~~~~~^
        kwargs[schema_arg.name], schema_arg.type
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    ),
    ^
  File "/.../lib/python3.14/site-packages/torch/_export/serde/serialize.py", line 1534, in serialize_input
    return Argument.create(as_scalar_type=_TORCH_TO_SERIALIZE_DTYPE[arg])
                                          ~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^
KeyError: torch.uint32

Versions

Collecting environment information...
PyTorch version: 2.12.0.dev20260402
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 26.3.1 (arm64)
GCC version: Could not collect
Clang version: 17.0.0 (clang-1700.6.3.2)
CMake version: version 4.2.3
Libc version: N/A

Python version: 3.14.2 (v3.14.2:df793163d58, Dec  5 2025, 12:18:06) [Clang 16.0.0 (clang-1600.0.26.6)] (64-bit runtime)
Python platform: macOS-26.3.1-arm64-arm-64bit-Mach-O
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
Is XPU available: False
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True
Caching allocator config: N/A

CPU:
Apple M4 Pro

Versions of relevant libraries:
[pip3] Could not collect
[conda] Could not collect

cc @chauhang @penguinwu @avikchaudhuri @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4

extent analysis

TL;DR

The most likely fix is to change the dtype of the tensor returned by the forward method to a supported type, such as torch.int32 or torch.float32.

Guidance

The error message indicates a KeyError exception when trying to serialize the tensor with dtype=torch.uint32. This suggests that torch.uint32 is not a supported dtype for serialization.
To fix this issue, you can try changing the dtype of the tensor returned by the forward method to a supported type, such as torch.int32 or torch.float32.
You can verify that the fix worked by running the code again and checking that the serialization is successful.
If you need to use torch.uint32 for your specific use case, you may need to consider alternative serialization methods or libraries that support this dtype.

Example

import torch

class Module(torch.nn.Module):
    def forward(self):
        return torch.empty((), dtype=torch.int32)

ep = torch.export.export(Module(), ())
torch.export.save(ep, "ep.pt2")

Notes

The PyTorch version used in this example is 2.12.0.dev20260402, which may not be the latest version. You should check the documentation for your specific version to see if torch.uint32 is supported.
The error message suggests that the issue is related to serialization, so you may need to investigate alternative serialization methods or libraries if you need to use torch.uint32.

Recommendation

Apply workaround: Change the dtype of the tensor returned by the forward method to a supported type, such as torch.int32 or torch.float32, as shown in the example above. This should allow the serialization to succeed.

Vote matrix · Quick signals

Works

Did the solution work? Tap to confirm.

Easy Fix

Was it a quick fix?

Time Saver

Did it save you time?

Blocking

Was it severely blocking?

Common Issue

Are others likely hitting this too?

Flaky / Intermittent

Is it intermittent?

Verified / Reproducible

Can you reproduce it reliably?

#indexing error #inference speed #output truncation #response parsing #generation error

Still need to ship something?

×6

Another batch ranked right after the header list — different links, same matching logic.

Data

Security

Network

Code

UI/UX

Text

System

Multimedia

Protocol

API

Engineering

pytorch - ✅(Solved) Fix `torch.export.save` fails when `torch.uint32` is present [1 pull requests, 1 comments, 1 participants]

Recommended Tools

GitHub issue graph ai analysis

Error Message

Fix Action

Fixed

PR fix notes

PR #179434: Support `torch.uint{32,64}` in `torch.export.save`

Description (problem / solution / changelog)

Changed files

Code Example

🐛 Describe the bug

Error logs

Versions

extent analysis

TL;DR

Guidance

Example

Notes

Recommendation

Still need to ship something?

TRENDING

pytorch - ✅(Solved) Fix `torch.export.save` fails when `torch.uint32` is present [1 pull requests, 1 comments, 1 participants]

Recommended Tools

GitHub issue graph ai analysis

Error Message

Fix Action

Fixed

PR fix notes

PR #179434: Support torch.uint{32,64} in torch.export.save

Description (problem / solution / changelog)

Changed files

Code Example

🐛 Describe the bug

Error logs

Versions

extent analysis

TL;DR

Guidance

Example

Notes

Recommendation

Still need to ship something?

RELATED_DISCOVERY

TRENDING

PR #179434: Support `torch.uint{32,64}` in `torch.export.save`