Fix imaginary part discarding warning in to function.

USTC-TNS · Nov 18, 2023 · 3658abe · 3658abe
1 parent a3a21f7
commit 3658abe
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/tat/tensor.py b/tat/tensor.py
@@ -422,12 +422,16 @@ def to(self: Tensor, new_type: typing.Any) -> Tensor:
                 new_type = torch.complex64
             elif new_type in ["complex128", "complex", "Z"]:
                 new_type = torch.complex128
+        if self.dtype.is_complex and new_type.is_floating_point:
+            data = self.data.real.to(new_type)
+        else:
+            data = self.data.to(new_type)
         return Tensor(
             names=self.names,
             edges=self.edges,
             fermion=self.fermion,
             dtypes=self.dtypes,
-            data=self.data.to(new_type),
+            data=data,
             mask=self.mask,
         )