allow the line profiler to work
MODIFY
test/test_conv_speed.py
test/test_conv_speed.py
#!/usr/bin/env python
-import time
-import cProfile
-import pstats
-import unittest
-from tinygrad.tensor import Tensor
+# if you'd like to use the line profiler
try:
import line_profiler
prof = line_profiler.LineProfiler()
@@ -14,6 +10,12 @@ try:
except ImportError:
prof = None
+import time
+import cProfile
+import pstats
+import unittest
+from tinygrad.tensor import Tensor
+
def profile_conv(bs, chans, conv, cnt=10):
img = Tensor.zeros(bs, 1, 28, 28)
conv = Tensor.randn(chans, 1, conv, conv)

allow the line profiler to work