Fixed code according to new D8 structure