oops optimizer now demands to know device_type

rovinyu · Feb 5, 2023 · 5809026 · 5809026
1 parent 34720df
commit 5809026
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/bench.py b/bench.py
@@ -57,7 +57,7 @@ def get_batch(split):
 model = GPT(gptconf)
 model.to(device)
 
-optimizer = model.configure_optimizers(weight_decay=1e-2, learning_rate=1e-4, betas=(0.9, 0.95))
+optimizer = model.configure_optimizers(weight_decay=1e-2, learning_rate=1e-4, betas=(0.9, 0.95), device_type=device_type)
 
 if compile:
     print("Compiling model...")