Display scores figure update comments

docs/source/quickstart.rst

@@ -337,7 +337,7 @@ saves the scores to disk in memory-mapped file (:code:`.mmap` format).
 We can visualize some of the top scoring :code:`TRAK` images from the
 :code:`scores` array we just computed:
 
-.. image:: assets/trak_scores_quickstart.png
+.. image:: ../assets/trak_scores_quickstart.png
    :alt: Top scoring TRAK images
 
 

trak/gradient_computers.py

@@ -111,7 +111,7 @@ def __init__(
         self.logger = logging.getLogger("GradientComputer")
 
     def load_model_params(self, model) -> None:
-        """Given a a torch.nn.Module model, inits/updates the (functional)
+        """Given a torch.nn.Module model, inits/updates the (functional)
         weights and buffers. See https://pytorch.org/docs/stable/func.html
         for more details on :code:`torch.func`'s functional models.
 

trak/savers.py

@@ -45,7 +45,7 @@ def __init__(
         TRAKer class has the same hyperparameters (metadata) as the one
         specified in the save directory. Next, this method loads any existing
         computed results / intermediate values in the save directory. Last, it
-        initalizes the self.current_store attributes which will be later
+        initializes the self.current_store attributes which will be later
         populated with data for the "current" model ID of the TRAKer instance.
 
         Args:
@@ -271,7 +271,7 @@ def register_model_id(
         """This method
         1) checks if the model ID already exists in the save dir
         2) if yes, it raises an error since model IDs must be unique
-        3) if not, it creates a metadata file for it and initalizes store mmaps
+        3) if not, it creates a metadata file for it and initializes store mmaps
 
         Args:
             model_id (int):

trak/traker.py

@@ -81,9 +81,9 @@ def __init__(
                 Size of the train set that TRAK is featurizing
             save_dir (str, optional):
                 Directory to save final TRAK scores, intermediate results, and
-                metadata. Defaults to :code:'./trak_results'.
+                metadata. Defaults to :code:`./trak_results`.
             load_from_save_dir (bool, optional):
-                If True, the :class`.TRAKer` instance will attempt to load
+                If True, the :class:`.TRAKer` instance will attempt to load
                 existing metadata from save_dir. May lead to I/O issues if
                 multiple TRAKer instances ran in parallel have this flag set to
                 True. See the SLURM tutorial for more details.