2 files changed, 11 insertions, 3 deletions
diff --git a/gpu.py b/gpu.py
index ed8ca12..e6b6856 100644
--- a/gpu.py
+++ b/gpu.py
@@ -259,7 +259,8 @@ class GPU(object):
                                       block=(self.nthread_per_block,1,1),
                                       grid=(blocks, 1))
 
-        #self.context.synchronize()
+        if 'profile' in __builtins__:
+            self.context.synchronize()
 
     def get_photons(self):
         '''Returns a dictionary of current photon state information.
@@ -311,7 +312,8 @@ class GPU(object):
                                                      self.earliest_time_gpu, 
                                                      block=(self.nthread_per_block,1,1), 
                                                      grid=(len(self.earliest_time_int_gpu)//self.nthread_per_block+1,1))
-        #self.context.synchronize()
+        if 'profile' in __builtins__:
+            self.context.synchronize()
         
 
     def get_hits(self):
diff --git a/sim.py b/sim.py
index 0d775f4..34477f9 100755
--- a/sim.py
+++ b/sim.py
@@ -75,7 +75,13 @@ def write_event(T, ev, event_id, hits, photon_start=None, photon_stop=None):
     root.fill_hits(ev, len(hits['t']), hits['t'], hits['q'], hits['history'])
     T.Fill()
 
-#@profile
+# Allow profile decorator to exist, but do nothing if not running under kernprof
+try:
+    profile = profile
+except NameError:
+    profile = lambda x: x
+
+@profile
 def main():
     parser = optparse.OptionParser('%prog')
     parser.add_option('-b', type='int', dest='nbits', default=10)