Weights not properly updating

Could you try to fix the weight init mentioned in your other thread?