update: test_scalable_shampoo_pre_conditioner_with_svd

kozistr · kozistr · commit 02aaecb526a1 · 2023-04-21T21:47:41.000+09:00
diff --git a/tests/test_optimizers.py b/tests/test_optimizers.py
@@ -316,7 +316,11 @@ def test_scalable_shampoo_pre_conditioner_with_svd(pre_conditioner_type):
     )
 
     optimizer = load_optimizer('scalableshampoo')(
-        model.parameters(), start_preconditioning_step=1, pre_conditioner_type=pre_conditioner_type, use_svd=True
+        model.parameters(),
+        start_preconditioning_step=1,
+        preconditioning_compute_steps=1,
+        pre_conditioner_type=pre_conditioner_type,
+        use_svd=True,
     )
     optimizer.zero_grad()
 

Original file line number	Diff line number	Diff line change
`@@ -316,7 +316,11 @@ def test_scalable_shampoo_pre_conditioner_with_svd(pre_conditioner_type):`
`316`	`316`	`)`
`317`	`317`
`318`	`318`	`optimizer = load_optimizer('scalableshampoo')(`
`319`		`- model.parameters(), start_preconditioning_step=1, pre_conditioner_type=pre_conditioner_type, use_svd=True`
	`319`	`+ model.parameters(),`
	`320`	`+ start_preconditioning_step=1,`
	`321`	`+ preconditioning_compute_steps=1,`
	`322`	`+ pre_conditioner_type=pre_conditioner_type,`
	`323`	`+ use_svd=True,`
`320`	`324`	`)`
`321`	`325`	`optimizer.zero_grad()`
`322`	`326`