small code fix in model