diff --git a/torchrec/sparse/jagged_tensor.py b/torchrec/sparse/jagged_tensor.py
index 6799ae9e0..9e353d47d 100644
--- a/torchrec/sparse/jagged_tensor.py
+++ b/torchrec/sparse/jagged_tensor.py
@@ -7,6 +7,7 @@
 
 import abc
 import operator
+import inspect
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import torch
@@ -635,7 +636,11 @@ def _jt_flatten_spec(t: JaggedTensor, spec: TreeSpec) -> List[Optional[torch.Ten
 
 
 register_pytree_node(
-    JaggedTensor, _jt_flatten, _jt_unflatten, flatten_with_keys_fn=_jt_flatten_with_keys
+    JaggedTensor,
+    _jt_flatten,
+    _jt_unflatten,
+    flatten_with_keys_fn=_jt_flatten_with_keys,
+    serialized_type_name="torchrec.JaggedTensor",
 )
 register_pytree_flatten_spec(JaggedTensor, _jt_flatten_spec)
 
@@ -2094,6 +2099,7 @@ def _kjt_flatten_spec(
     _kjt_flatten,
     _kjt_unflatten,
     flatten_with_keys_fn=_kjt_flatten_with_keys,
+    serialized_type_name="torchrec.KeyedJaggedTensor",
 )
 register_pytree_flatten_spec(KeyedJaggedTensor, _kjt_flatten_spec)