made test smaller

plutonium-239 · Apr 1, 2024 · 4c07ca1 · 4c07ca1
1 parent 1911520
commit 4c07ca1
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 5 deletions.
diff --git a/pyproject.toml b/pyproject.toml
@@ -6,7 +6,12 @@ requires = ["flit_core >=3.2,<4"]
 name = "memsave_torch"
 version = "1.0.0"
 dependencies = [
-  "torch"
+  "torch",
+  "torchvision",
+  "codetiming",
+  "memory_profiler",
+  "pandas",
+  "tqdm"
 ]
 requires-python = ">=3.8"
 authors = [

diff --git a/test/test_utils.py b/test/test_utils.py
@@ -29,11 +29,11 @@ def test_all():
     models = ["deeprelumodel"]
 
     models = prefix_in_pairs("memsave_", models)
-    batch_size = 64
+    batch_size = 2
     input_channels = 3
-    input_HW = 224
-    num_classes = 1000
-    device = "cuda"
+    input_HW = 32
+    num_classes = 20
+    device = "cpu"
     architecture = "conv"
     results_dir = ".test_results/"