diff --git a/preprocessing_distributions.py b/preprocessing_distributions.py
index 9a5ca64..4ce9ca0 100644
--- a/preprocessing_distributions.py
+++ b/preprocessing_distributions.py
@@ -43,7 +43,7 @@
                            dtype=None,
                            encoding=None)
 
-if len(restraints.shape) == 1:
+if len(restraints.shape) < 2:
     restraints = np.array([restraints])
 
 distogram = []