Enable profiling

Author: Joseph Suarez

pufferlib/pufferl.py

@@ -734,8 +734,8 @@ def __call__(self, name, epoch, nest=False):
         if (epoch + 1) % self.frequency != 0:
             return
 
-        #if torch.cuda.is_available():
-        #    torch.cuda.synchronize()
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
 
         tick = time.time()
         if len(self.stack) != 0 and not nest:
@@ -752,8 +752,8 @@ def pop(self, end):
         profile['elapsed'] += delta * self.frequency
 
     def end(self):
-        #if torch.cuda.is_available():
-        #    torch.cuda.synchronize()
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
 
         end = time.time()
         for i in range(len(self.stack)):