Lint

.pre-commit-config.yaml

@@ -10,7 +10,7 @@ repos:
     -   id: check-added-large-files
 -   repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.4.7
+    rev: v0.5.5
     hooks:
         # Run the linter.
         - id: ruff

bootstrap/factories.py

@@ -111,14 +111,14 @@ def parallelize_model(model: torch.nn.Module) -> torch.nn.Module:
 
 
 def make_optimizer(
-    optimizer_partial: Partial[torch.optim.Optimizer], model: torch.nn.Module
-) -> torch.optim.Optimizer:
+    optimizer_partial: Partial[torch.optim.optimizer.Optimizer], model: torch.nn.Module
+) -> torch.optim.optimizer.Optimizer:
     return optimizer_partial(model.parameters())
 
 
 def make_scheduler(
     scheduler_partial: Partial[torch.optim.lr_scheduler.LRScheduler],
-    optimizer: torch.optim.Optimizer,
+    optimizer: torch.optim.optimizer.Optimizer,
     epochs: int,
 ) -> torch.optim.lr_scheduler.LRScheduler:
     scheduler = scheduler_partial(

bootstrap/launch_experiment.py

@@ -108,7 +108,7 @@ def init_wandb(
 def launch_experiment(
     run,  # type: ignore
     data_loader: Partial[DataLoader[Any]],
-    optimizer: Partial[torch.optim.Optimizer],
+    optimizer: Partial[torch.optim.optimizer.Optimizer],
     scheduler: Partial[torch.optim.lr_scheduler.LRScheduler],
     trainer: Partial[BaseTrainer],
     tester: Partial[BaseTester],

conf/experiment.py

@@ -46,10 +46,11 @@
     zen_partial=True, populate_full_signature=False
 )
 
-""" ================== Dataset ================== """
+# ================== Dataset ==================
 
 
-# Dataclasses are a great and simple way to define a base config group with default values.
+# Dataclasses are a great and simple way to define a base config group with default
+# values.
 @dataclass
 class ExampleDatasetConf:
     dataset_name: str = "image_dataset"
@@ -86,7 +87,7 @@ class ExampleDatasetConf:
     name="image_a_tiny",
 )
 
-""" ================== Dataloader & sampler ================== """
+# ================== Dataloader & sampler ==================
 
 
 @dataclass
@@ -107,12 +108,12 @@ class DataloaderConf:
     persistent_workers: bool = False
 
 
-""" ================== Model ================== """
+# ================== Model ==================
 # Pre-set the group for store's model entries
 model_store = store(group="model")
 
-# Not that encoder_input_dim depend on dataset.img_dim, so we need to use a partial to set them in
-# the launch_experiment function.
+# Not that encoder_input_dim depend on dataset.img_dim, so we need to use a partial to
+# set them in the launch_experiment function.
 model_store(
     pbuilds(
         ExampleModel,
@@ -134,7 +135,7 @@ class DataloaderConf:
     name="model_b",
 )
 
-""" ================== Losses ================== """
+# ================== Losses ==================
 training_loss_store = store(group="training_loss")
 training_loss_store(
     pbuilds(
@@ -145,7 +146,7 @@ class DataloaderConf:
 )
 
 
-""" ================== Optimizer ================== """
+# ================== Optimizer ==================
 
 
 @dataclass
@@ -157,21 +158,21 @@ class Optimizer:
 opt_store = store(group="optimizer")
 opt_store(
     pbuilds(
-        torch.optim.Adam,
+        torch.optim.adam.Adam,
         builds_bases=(Optimizer,),
     ),
     name="adam",
 )
 opt_store(
     pbuilds(
-        torch.optim.SGD,
+        torch.optim.sgd.SGD,
         builds_bases=(Optimizer,),
     ),
     name="sgd",
 )
 
 
-""" ================== Scheduler ================== """
+# ================== Scheduler ==================
 sched_store = store(group="scheduler")
 sched_store(
     pbuilds(
@@ -197,7 +198,7 @@ class Optimizer:
     name="cosine",
 )
 
-""" ================== Experiment ================== """
+# ================== Experiment ==================
 
 
 @dataclass

src/base_trainer.py

@@ -23,7 +23,7 @@
 from rich.text import Text
 from torch import Tensor
 from torch.nn import Module
-from torch.optim import Optimizer
+from torch.optim.optimizer import Optimizer
 from torch.utils.data import DataLoader
 from torchmetrics import MeanMetric
 
@@ -54,7 +54,7 @@ def __init__(
         """Base trainer class.
         Args:
             model (torch.nn.Module): Model to train.
-            opt (torch.optim.Optimizer): Optimizer to use.
+            opt (torch.optim.optimizer.Optimizer): Optimizer to use.
             train_loader (torch.utils.data.DataLoader): Training dataloader.
             val_loader (torch.utils.data.DataLoader): Validation dataloader.
         """