Update README.md

This commit is contained in:
William Falcon 2020-05-12 08:46:55 -04:00 committed by GitHub
parent 39584d08ad
commit 423b82ea6c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 2 deletions

View File

@ -59,7 +59,7 @@ Here's a minimal example without a validation or test loop.
```python
# this is just a plain nn.Module with some structure
class MNISTModel(pl.LightningModule):
class LitClassifier(pl.LightningModule):
def __init__(self):
super(MNISTModel, self).__init__()
@ -80,7 +80,7 @@ class MNISTModel(pl.LightningModule):
# train!
train_loader = DataLoader(MNIST(os.getcwd(), train=True, download=True, transform=transforms.ToTensor()), batch_size=32)
mnist_model = MNISTModel()
mnist_model = LitClassifier()
trainer = pl.Trainer(gpus=8, precision=16)
trainer.fit(mnist_model, train_loader)
```