Skip to content
Snippets Groups Projects
Commit 70d11575 authored by William Falcon's avatar William Falcon
Browse files

clean up sample project

parent 041f25ad
No related branches found
No related tags found
No related merge requests found
from project.datasets.mnist import mnist
from project.lit_classifier_main import LitClassifier
from pytorch_lightning import Trainer, seed_everything
def test_lit_classifier():
seed_everything(1234)
model = LitClassifier()
train, val, test = mnist()
trainer = Trainer(limit_train_batches=50, limit_val_batches=20, max_epochs=2)
trainer.fit(model, train, val)
results = trainer.test(test_dataloaders=test)
assert results[0]['test_acc'] > 0.7
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment