dimension fix

AWS-BugBust-75 · Mar 9, 2021 · 542fafd · 542fafd
1 parent 698519d
commit 542fafd
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/deepchem/feat/complex_featurizers/rdkit_grid_featurizer.py b/deepchem/feat/complex_featurizers/rdkit_grid_featurizer.py
@@ -449,6 +449,7 @@ def _featurize(self, complex):
         else:
           features_dict[system_id] = np.concatenate(feature_arrays, axis=-1)
 
-    # TODO(rbharath): Is this squeeze OK?
-    features = np.squeeze(np.array(list(features_dict.values())))
+    features = np.array(list(features_dict.values()))
+    if self.nb_rotations == 0:  # squeeze out axis with dimension 1
+      features = np.squeeze(features, axis=0)
     return features
diff --git a/deepchem/feat/tests/test_rdkit_grid_features.py b/deepchem/feat/tests/test_rdkit_grid_features.py
@@ -134,6 +134,8 @@ def test_custom_cutoffs(self):
   def test_rotations(self):
     featurizer = RdkitGridFeaturizer(
         nb_rotations=3,
+        box_width=16.,
+        voxel_width=1.,
         feature_types=['voxel_combined'],
         flatten=False,
         sanitize=True)