Add stage marker to test_memory (#820)

pyro-ppl · Feb 26, 2018 · 6df6721 · 6df6721
1 parent 0f9b498
commit 6df6721
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/tests/test_memory.py b/tests/test_memory.py
@@ -3,6 +3,7 @@
 import gc
 
 import networkx as nx
+import pytest
 import torch
 from torch.autograd import Variable
 
@@ -14,6 +15,9 @@
 from pyro.poutine.trace import Trace
 
 
+pytestmark = pytest.mark.stage('unit')
+
+
 def count_objects_of_type(type_):
     return sum(1 for obj in gc.get_objects() if isinstance(obj, type_))