From 03fc6e8a49c1b358eb89c7eafbf4dbe588882621 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Wed, 21 Oct 2020 08:48:16 -0700 Subject: [PATCH] layer init for all sizes --- tinygrad/utils.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tinygrad/utils.py b/tinygrad/utils.py index 31791b90eb..2887748d62 100644 --- a/tinygrad/utils.py +++ b/tinygrad/utils.py @@ -1,7 +1,7 @@ import numpy as np -def layer_init_uniform(m, h): - ret = np.random.uniform(-1., 1., size=(m,h))/np.sqrt(m*h) +def layer_init_uniform(*x): + ret = np.random.uniform(-1., 1., size=x)/np.sqrt(np.prod(x)) return ret.astype(np.float32) def fetch_mnist():