Small modifications / fix