Some small code updates