ml: docstrings fixed

miniff/ml.py

@@ -913,7 +913,7 @@ class Normalization:
         Returns
         -------
         result : torch.Tensor
-            The resulting counts.
+            A 2D tensor `[n_samples, len(per_point_datasets)]` with counts.
         """
         return torch.cat(tuple(
             i.mask.sum(dim=1)[:, None]
@@ -1088,10 +1088,10 @@ class Normalization:
         offset_features : bool
         offset_charges : bool
             If True, offsets energies, descriptors, and/or charges.
-        scale_energy : bool
-        scale_features : bool
-        scale_charges : bool
-        scale_energy_gradients : bool
+        scale_energy : float
+        scale_features : float
+        scale_charges : float
+        scale_energy_gradients : float
             If set scales energies, descriptors, and/or charges to the value specified.
 
         Returns

miniff/ml_util.py

@@ -468,7 +468,7 @@ class SimpleClosure:
             A function `loss(networks, data, criterion, **loss_kwargs)`
             returning `loss_result` tuple.
         dataset : Dataset
-            Default dataset to compute loss for.
+            Default dataset to compute the loss for.
         criterion : torch.nn.Module
             Loss criterion, defaults to MSE.
         optimizer : torch.optim.Optimizer