From 70d11575aef4c73f95b542c57a71cffb688aab2c Mon Sep 17 00:00:00 2001
From: William Falcon <waf2107@columbia.edu>
Date: Fri, 11 Sep 2020 11:04:54 -0400
Subject: [PATCH] clean up sample project

---
 tests/__init__.py        |  0
 tests/test_classifier.py | 15 +++++++++++++++
 2 files changed, 15 insertions(+)
 create mode 100644 tests/__init__.py
 create mode 100644 tests/test_classifier.py

diff --git a/tests/__init__.py b/tests/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/test_classifier.py b/tests/test_classifier.py
new file mode 100644
index 0000000..a11a314
--- /dev/null
+++ b/tests/test_classifier.py
@@ -0,0 +1,15 @@
+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
-- 
GitLab