Add stage marker to test_memory
@@ -3,6 +3,7 @@ from __future__ import absolute_import, division, print_function
import gc
import networkx as nx
+import pytest
import torch
from torch.autograd import Variable
@@ -14,6 +15,9 @@ from pyro.optim import Adam
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_))
