diff --git a/velociraptor/observations/objects.py b/velociraptor/observations/objects.py
index efb9973..601446f 100644
--- a/velociraptor/observations/objects.py
+++ b/velociraptor/observations/objects.py
@@ -13,7 +13,7 @@
 from matplotlib import rcParams
 
 from astropy.units import Quantity
-from astropy.cosmology.core import Cosmology
+from astropy.cosmology import Cosmology
 from astropy.cosmology import wCDM, FlatLambdaCDM
 
 import h5py
diff --git a/velociraptor/units.py b/velociraptor/units.py
index a66ac64..2c6a63a 100644
--- a/velociraptor/units.py
+++ b/velociraptor/units.py
@@ -7,7 +7,7 @@
 import unyt
 import h5py
 
-from astropy.cosmology.core import Cosmology
+from astropy.cosmology import Cosmology
 from astropy.cosmology import wCDM, FlatLambdaCDM