Update docs

docs/source/index.rst

@@ -80,20 +80,20 @@ TODO
 
         Enable users to make predictions from the command line.
 
-    .. card:: 2️⃣  Support more architectures
+    .. card:: 2️⃣  Add pretrained models
         :text-align: center
 
-        This might require the usage of |pyg|.
+        Enable users to fine-tune models trained on large data.
 
-    .. card:: 3️⃣  Extend featurization
+    .. card:: 3️⃣  Support more architectures
         :text-align: center
 
-        Add more featurization options. These should be fast!
+        This might require the usage of |pyg|.
 
-    .. card:: 4️⃣  Add pretrained models
+    .. card:: 4️⃣  Extend featurization
         :text-align: center
 
-        Enable users to fine-tune models trained on large data.
+        Add more featurization options. These should be fast!
 
 Contributing
 ------------
@@ -117,10 +117,6 @@ discussions.
 Citing
 ------
 
-.. admonition:: Citation
-
-    Here is the text for citation.
-
 If you use AIdosrb in your research, please consider citing the following work::
 
 	Currently, not available.

src/aidsorb/datamodules.py

@@ -45,7 +45,8 @@ class PCDDataModule(L.LightningDataModule):
         directory structure with :func:`~aidsorb.data.prepare_data`.
 
     .. todo::
-        Add support for ``predict_dataloader``.
+        * Add support for ``predict_dataloader``.
+        * Add option ``drop_last`` for ``train_dataloader``.
 
     Parameters
     ----------

src/aidsorb/litmodels.py

@@ -26,7 +26,7 @@
 
 class PointLit(L.LightningModule):
     r"""
-    :class:`~lightning.LightningModule` for :class:`aidsorb.models`.
+    ``LightningModule`` for :class:`aidsorb.models`.
 
     .. _loss functions: https://pytorch.org/docs/stable/nn.html#loss-functions
     .. _monitor: https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.callbacks.ModelCheckpoint.html#modelcheckpoint
@@ -107,9 +107,9 @@ def training_step(self, batch, batch_idx):
         Also, make predictions that will be used on epoch-level operations.
 
         .. note::
-            The ``BatchNorm`` and ``Dropout`` are set in inference mode during
-            predictions, so an accurate estimate of training performance is
-            reported.
+            Inference mode is enabled during predictions, so an accurate
+            estimate of training performance (e.g. when using
+            :class:`~torch.nn.Dropout`) is reported.
         """
         assert self.training
         assert torch.is_grad_enabled()

tests/test_utils.py

@@ -52,7 +52,7 @@ def test_features(self):
                 )
 
         self.assertEqual(name, 'H2O')
-        self.assertTrue(np.all(pcd == water))
+        self.assertTrue(np.array_equal(pcd, water))
 
 
 class TestPCDFromFiles(unittest.TestCase):