fixed a small bug in 2x2 case of initialize weights