From 0bbf66627cbb43534550c05aec78998cf20fad2b Mon Sep 17 00:00:00 2001
From: Mufeed VH <mufeedvh@gmail.com>
Date: Wed, 25 Nov 2020 09:09:13 +0530
Subject: [PATCH] Define `ProfileOp` class once (#131)

* define `ProfileOp` class once

* clean `ProfileOp` class

* removed `else: pass`
---
 tinygrad/tensor.py | 24 +++++++++---------------
 1 file changed, 9 insertions(+), 15 deletions(-)

diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py
index 36aa458c1c..e3621e2e84 100644
--- a/tinygrad/tensor.py
+++ b/tinygrad/tensor.py
@@ -19,25 +19,19 @@ if DEBUG:
     for name, _ in sorted(debug_times.items(), key=lambda x: -x[1]):
       print("%20s : %3d  %10.2f ms" % (name, debug_counts[name], debug_times[name]))
   atexit.register(print_debug_exit)
-  class ProfileOp:
-    def __init__(self, name, x, backward=False):
-      self.name = ("back_" if backward else "")+name
-      self.x = x
-    def __enter__(self):
-      self.st = time.time()
-    def __exit__(self, *junk):
+
+class ProfileOp:
+  def __init__(self, name, x, backward=False):
+    self.name = ("back_" if backward else "")+name
+    self.x = x
+  def __enter__(self):
+    if DEBUG: self.st = time.time()
+  def __exit__(self, *junk):
+    if DEBUG:
       et = (time.time()-self.st)*1000.
       debug_counts[self.name] += 1
       debug_times[self.name] += et
       print("%20s : %7.2f ms  %s" % (self.name, et, [y.shape for y in self.x]))
-else:
-  class ProfileOp:
-    def __init__(self, name, x, backward=False):
-      pass
-    def __enter__(self):
-      pass
-    def __exit__(self, *junk):
-      pass
 
 cl_ctx, cl_queue = None, None
 def require_init_gpu():