Add a thumbnail (image shows up now!)

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

examples/tutorials/plot_5_warm_starting.py

@@ -118,18 +118,16 @@ class Net(nn.Sequential):
 
     def __init__(self, n_classes: int = 10):
         super().__init__(
-            nn.LazyConv2d(
-                32, 3, 1
-            ),  # NOTE: `in_channels` is determined in the first forward pass
+            # NOTE: `in_channels` is determined in the first forward pass
+            nn.LazyConv2d(32, 3, 1),
             nn.ReLU(),
             nn.Conv2d(32, 64, 3, 1),
             nn.ReLU(),
             nn.MaxPool2d(2),
             nn.Dropout2d(0.25),
             nn.Flatten(),
-            nn.LazyLinear(
-                128
-            ),  # NOTE: `in_features` is determined in the first forward pass
+            # NOTE: `in_features` is determined in the first forward pass
+            nn.LazyLinear(128),
             nn.ReLU(),
             nn.Dropout(0.5),
             nn.Linear(128, n_classes),
@@ -174,12 +172,10 @@ def test_epoch(model: Net, device: torch.device, test_loader: DataLoader) -> flo
         for data, target in test_loader:
             data, target = data.to(device), target.to(device)
             output = model(data)
-            test_loss += F.nll_loss(
-                output, target, reduction="sum"
-            ).item()  # sum up batch loss
-            pred = output.argmax(
-                dim=1, keepdim=True
-            )  # get the index of the max log-probability
+            # sum up batch loss
+            test_loss += F.nll_loss(output, target, reduction="sum").item()
+            # get the index of the max log-probability
+            pred = output.argmax(dim=1, keepdim=True)
             correct += pred.eq(target.view_as(pred)).sum().item()
 
     test_loss /= num_batches
@@ -379,4 +375,7 @@ def main(**kwargs):
     },
 )
 fig.show()
+fig.write_image("../../docs/src/_static/warm_start_thumbnail.png")
 fig
+
+# sphinx_gallery_thumbnail_path = '_static/warm_start_thumbnail.png'