diff --git a/models/rfd3/src/rfd3/model/inference_sampler.py b/models/rfd3/src/rfd3/model/inference_sampler.py
index 2c1482f..1e17dc7 100644
--- a/models/rfd3/src/rfd3/model/inference_sampler.py
+++ b/models/rfd3/src/rfd3/model/inference_sampler.py
@@ -13,6 +13,9 @@
     uniform_random_rotation,
 )
 
+from rfd3.inference.symmetry.atom_array import FIXED_ENTITY_ID, FIXED_TRANSFORM_ID
+from rfd3.inference.symmetry.symmetry_utils import apply_symmetry_to_xyz_atomwise
+
 ranked_logger = RankedLogger(__name__, rank_zero_only=True)