From a2d81514fd22c02f9322985e51a2bb197929506d Mon Sep 17 00:00:00 2001 From: Sohaib Date: Tue, 30 Apr 2024 00:26:14 -0700 Subject: [PATCH] just get dtype from kwargs (#4355) --- tinygrad/tensor.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index d123b64989..4d6043ef9e 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -514,7 +514,8 @@ class Tensor: t = Tensor.randint(2, 3, low=5, high=10) print(t.numpy()) """ - return Tensor.uniform(*shape, low=low, high=high, dtype=dtypes.int32, **kwargs) + assert dtypes.is_int(dtype := kwargs.pop("dtype", dtypes.int32)), f"Unsupported dtype {dtype} for randint" + return Tensor.uniform(*shape, low=low, high=high, dtype=dtype, **kwargs) @staticmethod def normal(*shape, mean=0.0, std=1.0, **kwargs) -> Tensor: