refactor

Author: xadupre

_doc/api/torch_export_patches/index.rst

@@ -8,6 +8,7 @@ onnx_diagnostic.torch_export_patches
     eval/index
     onnx_export_errors
     onnx_export_serialization
+    onnx_export_serialization_impl
     patches/index
     patch_expressions
     patch_inputs

_doc/api/torch_export_patches/onnx_export_serialization_impl.rst

@@ -0,0 +1,7 @@
+
+onnx_diagnostic.torch_export_patches.onnx_export_serialization_impl
+===================================================================
+
+.. automodule:: onnx_diagnostic.torch_export_patches.onnx_export_serialization_impl
+    :members:
+    :no-undoc-members:

_unittests/ut_torch_export_patches/test_patch_serialization.py

@@ -224,6 +224,13 @@ def test_unet_2d_condition_output(self):
 
         bo = UNet2DConditionOutput(sample=torch.rand((4, 4, 4)))
         self.assertEqual(bo.__class__.__name__, "UNet2DConditionOutput")
+        bo2 = torch_deepcopy([bo])
+        self.assertIsInstance(bo2, list)
+        self.assertEqual(
+            "UNet2DConditionOutput(sample:T1s4x4x4)",
+            self.string_type(bo, with_shape=True),
+        )
+
         with torch_export_patches():
             # internal function
             bo2 = torch_deepcopy([bo])