Update weights_handler.py