Some more code fix