Merge "[ML2] Small refactoring"