Re-added a missing import