Adding atom_features_size property to CommonArgs

chemprop/args.py

@@ -100,6 +100,7 @@ class CommonArgs(Tap):
 
     def __init__(self, *args, **kwargs):
         super(CommonArgs, self).__init__(*args, **kwargs)
+        self._atom_features_size = 0
         self._atom_descriptors_size = 0
 
     @property
@@ -129,6 +130,15 @@ def features_scaling(self) -> bool:
         """Whether to apply normalization with a :class:`~chemprop.data.scaler.StandardScaler` to the additional molecule-level features."""
         return not self.no_features_scaling
 
+    @property
+    def atom_features_size(self) -> int:
+        """The size of the atom features."""
+        return self._atom_features_size
+
+    @atom_features_size.setter
+    def atom_features_size(self, atom_features_size: int) -> None:
+        self._atom_features_size = atom_features_size
+
     @property
     def atom_descriptors_size(self) -> int:
         """The size of the atom descriptors."""

chemprop/train/cross_validate.py

@@ -68,13 +68,13 @@ def cross_validate(args: TrainArgs,
     )
     validate_dataset_type(data, dataset_type=args.dataset_type)
     args.features_size = data.features_size()
-    args.atom_features_size = 0
 
     if args.atom_descriptors == 'descriptor':
         args.atom_descriptors_size = data.atom_descriptors_size()
         args.ffn_hidden_size += args.atom_descriptors_size
     elif args.atom_descriptors == 'feature':
-        set_extra_atom_fdim(data.atom_features_size())
+        args.atom_features_size = data.atom_features_size()
+        set_extra_atom_fdim(args.atom_features_size)
 
     debug(f'Number of tasks = {args.num_tasks}')