Fix import errors

Fix #1