diff --git a/py4DSTEM/process/classification/featurization.py b/py4DSTEM/process/classification/featurization.py
index 26a63a62d..1f7d493ef 100644
--- a/py4DSTEM/process/classification/featurization.py
+++ b/py4DSTEM/process/classification/featurization.py
@@ -477,7 +477,8 @@ def get_class_DPs(self, datacube, method, thresh):
                     datacube.data.shape[2],
                     datacube.data.shape[3],
                 )
-            except ValueError:
+            # TODO Work out if its a ValueError
+            except Exception:
                 raise ValueError(
                     "Datacube must have same R_Nx and R_Ny dimensions as Featurization instance."
                 )