From c3d81bba2a08d5dae7610ef0bdb56f3dc6101a5e Mon Sep 17 00:00:00 2001 From: George Hotz Date: Mon, 6 Feb 2023 08:55:41 -0600 Subject: [PATCH] test_train: Adam -> SGD --- test/test_train.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/test_train.py b/test/test_train.py index 149bfa712..4aa36ff03 100644 --- a/test/test_train.py +++ b/test/test_train.py @@ -18,7 +18,7 @@ def train_one_step(model,X,Y): pcount = 0 for p in params: pcount += np.prod(p.shape) - optimizer = optim.Adam(params, lr=0.001) + optimizer = optim.SGD(params, lr=0.001) print("stepping %r with %.1fM params bs %d" % (type(model), pcount/1e6, BS)) st = time.time() train(model, X, Y, optimizer, steps=1, BS=BS)