zhangyiss/mlx-examples
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx-examples.git synced 2025-08-30 02:53:41 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

default to fp32 for now

Browse Source
This commit is contained in:
Awni Hannun 2023-12-18 17:15:49 -08:00
parent 05a8464d78
commit 688795c665
1 changed files with 1 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
t5/t5.py
View File

@ -146,9 +146,6 @@ class RMSNorm(nn.Module):
def __call__(self, x): def __call__(self, x):
t = x.dtype t = x.dtype
if t == mx.float16:
x = x.astype(mx.float32)
x = mx.clip(x, a_min=-1e9, a_max=1e9)
output = self._norm(x).astype(t) output = self._norm(x).astype(t)
return self.weight * output return self.weight * output
@ -406,7 +403,7 @@ if __name__ == "__main__":
help="The model data type.", help="The model data type.",
type=str, type=str,
choices=["float16", "bfloat16", "float32"], choices=["float16", "bfloat16", "float32"],
default="float16", default="float32",
) )
parser.add_argument("--seed", type=int, default=0, help="The PRNG seed") parser.add_argument("--seed", type=int, default=0, help="The PRNG seed")