fix: docstring and kwargs (#59)

pytorch-ignite · Apr 10, 2021 · fefaa33 · fefaa33
1 parent 818585e
commit fefaa33
Show file tree

Hide file tree

Showing 4 changed files with 77 additions and 38 deletions.
diff --git a/templates/gan/trainers.py b/templates/gan/trainers.py
@@ -34,17 +34,28 @@ def train_function(
 
     Parameters
     ----------
-    - config: config object
-    - engine: Engine instance
-    - batch: batch in current iteration
-    - netD: discriminator model
-    - netG: generator model
-    - loss_fn: nn.Module loss
-    - optimizerD: discriminator optimizer
-    - optimizerG: generator optimizer
-    - device: device to use for training
-    - real_labels: real label tensor
-    - fake_labels: fake label tensor
+    config
+        config object
+    engine
+        Engine instance
+    batch
+        batch in current iteration
+    netD
+        discriminator model
+    netG
+        generator model
+    loss_fn
+        nn.Module loss
+    optimizerD
+        discriminator optimizer
+    optimizerG
+        generator optimizer
+    device
+        device to use for training
+    real_labels
+        real label tensor
+    fake_labels
+        fake label tensor
 
     Returns
     -------

diff --git a/templates/gan/utils.py b/templates/gan/utils.py
@@ -37,8 +37,10 @@ def initialize(
 
     Parameters
     ----------
-    config:
+    config
         config object
+    num_channels
+        number of channels for Generator
 
     Returns
     -------

diff --git a/templates/image_classification/trainers.py b/templates/image_classification/trainers.py
@@ -31,13 +31,20 @@ def train_function(
 
     Parameters
     ----------
-    - config: config object
-    - engine: Engine instance
-    - batch: batch in current iteration
-    - model: nn.Module model
-    - loss_fn: nn.Module loss
-    - optimizer: torch optimizer
-    - device: device to use for training
+    config
+        config object
+    engine
+        Engine instance
+    batch
+        batch in current iteration
+    model
+        nn.Module model
+    loss_fn
+        nn.Module loss
+    optimizer
+        torch optimizer
+    device
+        device to use for training
 
     Returns
     -------
@@ -86,12 +93,16 @@ def evaluate_function(
 
     Parameters
     ----------
-    - config: config object
-    - engine: Engine instance
-    - batch: batch in current iteration
-    - model: nn.Module model
-    - loss_fn: nn.Module loss
-    - device: device to use for training
+    config
+        config object
+    engine
+        Engine instance
+    batch
+        batch in current iteration
+    model
+        nn.Module model
+    device
+        device to use for training
 
     Returns
     -------
@@ -129,6 +140,7 @@ def create_trainers(**kwargs) -> Tuple[Engine, Engine]:
             **kwargs,
         )
     )
+    kwargs.pop('optimizer')
     eval_engine = Engine(
         lambda e, b: evaluate_function(
             engine=e,

diff --git a/templates/single/trainers.py b/templates/single/trainers.py
@@ -31,13 +31,20 @@ def train_function(
 
     Parameters
     ----------
-    - config: config object
-    - engine: Engine instance
-    - batch: batch in current iteration
-    - model: nn.Module model
-    - loss_fn: nn.Module loss
-    - optimizer: torch optimizer
-    - device: device to use for training
+    config
+        config object
+    engine
+        Engine instance
+    batch
+        batch in current iteration
+    model
+        nn.Module model
+    loss_fn
+        nn.Module loss
+    optimizer
+        torch optimizer
+    device
+        device to use for training
 
     Returns
     -------
@@ -87,12 +94,18 @@ def evaluate_function(
 
     Parameters
     ----------
-    - config: config object
-    - engine: Engine instance
-    - batch: batch in current iteration
-    - model: nn.Module model
-    - loss_fn: nn.Module loss
-    - device: device to use for training
+    config
+        config object
+    engine
+        Engine instance
+    batch
+        batch in current iteration
+    model
+        nn.Module model
+    loss_fn
+        nn.Module loss
+    device
+        device to use for training
 
     Returns
     -------
@@ -133,6 +146,7 @@ def create_trainers(**kwargs) -> Tuple[Engine, Engine]:
             **kwargs,
         )
     )
+    kwargs.pop('optimizer')
     eval_engine = Engine(
         lambda e, b: evaluate_function(
             engine=e,