 (defn train
   []
   (let [network (create-network)
-        optimizer (create-sgd-optimizer network)
+        optimizer (create-optimizer network)
         epoch-batches (repeatedly n-epochs
                        #(into [] (partition batch-size (shuffle (range (count training-data))))))
         [optimizer network] (reduce (fn [opt-network batch-index-seq]
                                       (reduce (fn [[optimizer network] batch-indexes]
                                                 (let [input-seq (mapv training-data batch-indexes)
-                                                      answer-seq (mapv training-labels batch-indexes)]
-                                                  (train-batch input-seq answer-seq network optimizer loss-fn)))
+                                                      answer-seq (mapv training-labels batch-indexes)
+                                                      [optimizer network] (train-batch input-seq
+                                                                                       answer-seq
+                                                                                       network optimizer loss-fn)]
+                                                  ;(println "loss after batch:" (:loss network))
+                                                  [optimizer network]))
                                               opt-network
                                               batch-index-seq))
                                     [optimizer network]

(def hidden-layer-size 30)


(def n-epochs 1)
(def learning-rate 0.001)
(def momentum 0.9)
(def batch-size 100)
(def loss-fn (opt/mse-loss))

(defn random-matrix

                          (train-step input answer network loss-fn))
                        network
                        (map vector input-seq label-seq))
        batch-scale (/ 10.0 (count input-seq))]
    (m/scale! (core/gradient network) batch-scale)
    (core/optimise optimizer network)))

