Fix weight scaling for first layer.

Fixes #3.
Fix momentum bug.

Fixes #4.