Minor changes to docstrings for docs

Signed-off-by: Fabrice Normandin <[email protected]>

docs/macros.py

@@ -8,9 +8,23 @@
 from pathlib import Path
 from typing import Any
 
+import torch
+
 if typing.TYPE_CHECKING:
     from mkdocs_macros.plugin import MacrosPlugin
 
+import lightning
+from mkdocs_autoref_plugin.autoref_plugin import default_reference_sources
+
+default_reference_sources.extend(
+    [
+        lightning.Trainer,
+        lightning.LightningModule,
+        lightning.LightningDataModule,
+        torch.nn.Module,
+    ]
+)
+
 logger = logging.getLogger(__name__)
 
 

project/datamodules/__init__.py

@@ -1,6 +1,6 @@
 """Datamodules (datasets + preprocessing + dataloading)
 
-See the :ref:`lightning.LightningDataModule` class for more information.
+See the `lightning.LightningDataModule` class for more information.
 """
 
 from .image_classification import ImageClassificationDataModule

project/datamodules/image_classification/imagenet.py

@@ -46,10 +46,10 @@ class ImageNetDataModule(ImageClassificationDataModule):
     - Made this a subclass of VisionDataModule
 
     Notes:
+
     - train_dataloader uses the train split of imagenet2012 and puts away a portion of it for the validation split.
     - val_dataloader uses the part of the train split of imagenet2012  that was not used for training via
         `num_imgs_per_val_class`
-        - TODO: needs to pass split='val' to UnlabeledImagenet.
     - test_dataloader uses the validation split of imagenet2012 for testing.
         - TODO: need to pass num_imgs_per_class=-1 for test dataset and split="test".
     """