merge

1 files changed, 6 insertions, 2 deletions

diff --git a/tests/test_pdf.py b/tests/test_pdf.py
index c7c9394..571cbd4 100644
--- a/tests/test_pdf.py
+++ b/tests/test_pdf.py
@@ -11,7 +11,7 @@ from chroma.sim import Simulation
 
 class TestPDF(unittest.TestCase):
     def setUp(self):
-        self.detector = chroma.detectors.find("microlbne")
+        self.detector = chroma.detectors.microlbne()
         self.detector.build()
         self.vertex_gen = constant_particle_gun('e-', (0,0,0), (1,0,0), 10)
 
@@ -19,7 +19,9 @@ class TestPDF(unittest.TestCase):
         '''Create a hit count and (q,t) PDF for 10 MeV events in MicroLBNE'''
 
         g4generator = G4ParallelGenerator(1, water_wcsim)
-        gpu_instance = gpu.GPU(0)
+
+        context = gpu.create_cuda_context()
+
         gpu_geometry = gpu.GPUGeometry(self.detector)
 
         nthreads_per_block, max_blocks = 64, 1024
@@ -48,6 +50,8 @@ class TestPDF(unittest.TestCase):
         for i, nhits in enumerate(hitcount):
             self.assertEqual(nhits, pdf[i].sum())
 
+        context.pop()
+
     def testSimPDF(self):
         sim = Simulation(self.detector)