Fix model weights checks