diff --git a/tests/gflownet/envs/test_lattice_parameters.py b/tests/gflownet/envs/test_lattice_parameters.py
index 94f212e72..e2dfb0398 100644
--- a/tests/gflownet/envs/test_lattice_parameters.py
+++ b/tests/gflownet/envs/test_lattice_parameters.py
@@ -9,19 +9,18 @@
 import pytest
 import torch
 
-from gflownet.envs.crystals.lattice_parameters import (
+from gflownet.envs.crystals.lattice_parameters import PARAMETER_NAMES, LatticeParameters
+from gflownet.utils.common import tfloat
+from gflownet.utils.crystals.constants import (
     CUBIC,
     HEXAGONAL,
     LATTICE_SYSTEMS,
     MONOCLINIC,
     ORTHORHOMBIC,
-    PARAMETER_NAMES,
     RHOMBOHEDRAL,
     TETRAGONAL,
     TRICLINIC,
-    LatticeParameters,
 )
-from gflownet.utils.common import tfloat
 
 N_REPETITIONS = 100