set dtype

Author: danielvegamyhre

torchao/prototype/moe_training/tensor.py

@@ -64,7 +64,7 @@ def __init__(
         tensor: torch.Tensor,
         dtype: torch.dtype,
     ):
-        self._data = tensor
+        self._data = tensor.to(dtype)
         self._dtype = dtype
 
     @classmethod