diff --git a/src/braket/ahs/atom_arrangement.py b/src/braket/ahs/atom_arrangement.py
index 381fe36b8..3d486c23c 100644
--- a/src/braket/ahs/atom_arrangement.py
+++ b/src/braket/ahs/atom_arrangement.py
@@ -32,6 +32,7 @@
 from numbers import Number
 from typing import Union, Tuple, List
 import numpy as np
+from braket.ahs.discretization_types import DiscretizationError, DiscretizationProperties
 
 
 class SiteType(Enum):