fix type in PHETensor

Signed-off-by: sagewe <wbwmat@gmail.com>
FederatedAI · Aug 22, 2023 · df770b9 · df770b9
1 parent 45afe68
commit df770b9
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/python/fate/arch/tensor/phe/_tensor.py b/python/fate/arch/tensor/phe/_tensor.py
@@ -35,8 +35,8 @@ def __init__(self, pk, evaluator, coder, shape: torch.Size, data, dtype, device)
         self._dtype = dtype
         self._device = device
 
-    def type(self):
-        return f"phe.{self.dtype}"
+    def type(self, dtype):
+        return self.with_template(self._data, dtype)
 
     def __repr__(self) -> str:
         return f"<PHETensor shape={self.shape}, dtype={self.dtype}>"