diff options
Diffstat (limited to 'examples/mnist/mnist.py')
| -rw-r--r-- | examples/mnist/mnist.py | 5 |
1 files changed, 2 insertions, 3 deletions
diff --git a/examples/mnist/mnist.py b/examples/mnist/mnist.py index 0a81878..a1706be 100644 --- a/examples/mnist/mnist.py +++ b/examples/mnist/mnist.py @@ -24,7 +24,7 @@ def train_model(name: str, dim): optimizer = torch.optim.Adam(net.parameters(), lr=0.5e-4) print(f"Training {name} ({dim} neurons)...") - for epoch in range(100): + for epoch in range(10): global loss for data in trainloader: inputs, targets = data @@ -33,8 +33,7 @@ def train_model(name: str, dim): loss = loss_fn(outputs, targets) loss.backward() optimizer.step() - if (epoch + 1) % 10 == 0: - print(f" Epoch {epoch+1}, Loss: {loss.item():.4f}") + print(f" Epoch {epoch+1}, Loss: {loss.item():.4f}") return net if __name__ == "__main__": |