minor optimization nits

felipecsl · May 3, 2024 · 3f33b2b · 3f33b2b
1 parent ddc8720
commit 3f33b2b
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 4 deletions.
diff --git a/.idea/gradle.xml b/.idea/gradle.xml
diff --git a/src/test/kotlin/OptimizationTest.kt b/src/test/kotlin/OptimizationTest.kt
@@ -24,19 +24,21 @@ class OptimizationTest {
 
     // gradient descent
     val maxSteps = 1000
-    for (k in 0.rangeTo(maxSteps)) {
+    for (k in 0..maxSteps) {
       // forward
       val (totalLoss, predictions) = loss()
       // backward
       model.zeroGrad()
       totalLoss.backward()
       // update(sgd)
-      val learningRate = 0.01//1.0 - 0.9 * k / 100
+      val learningRate = 0.02
       for (p in model.parameters()) {
         p.data -= learningRate * p.grad
       }
-      println("step $k loss ${totalLoss.data}")
-      if (k == maxSteps - 1) {
+      if (k % 100 == 0) {
+        println("step $k loss ${totalLoss.data}, learningRate=$learningRate")
+      }
+      if (k == maxSteps) {
         println("predictions: ${predictions.map { it.data }}")
       }
     }