Fixed minor bugs

Fixed minor bugs in trainWeight module