fix: do not serialize to list on model dump

Author: samedii

lantern/functional_base.py

@@ -32,3 +32,18 @@ def staticmethod(cls, fn):
     @classmethod
     def classmethod(cls, fn):
         return cls.setattr(fn.__name__, classmethod(fn))
+
+
+def test_replace_same_device():
+    import torch
+
+    from .tensor import Tensor
+
+    class A(FunctionalBase):
+        x: Tensor
+        y: int
+
+    a = A(x=torch.tensor([1, 2, 3]).to("meta"), y=2)
+    b = a.replace(y=2)
+
+    assert b.x.device == a.x.device

lantern/numpy.py

@@ -54,7 +54,7 @@ def __get_pydantic_core_schema__(
                 ]
             ),
             serialization=core_schema.plain_serializer_function_ser_schema(
-                lambda instance: instance.tolist()
+                lambda instance: instance
             ),
         )
 

lantern/tensor.py

@@ -54,7 +54,7 @@ def __get_pydantic_core_schema__(
                 ]
             ),
             serialization=core_schema.plain_serializer_function_ser_schema(
-                lambda instance: instance.tolist()
+                lambda instance: instance
             ),
         )