Update

Deep-Learning-Profiling-Tools · Mar 27, 2024 · f198419 · f198419
1 parent d43bd5d
commit f198419
Showing 1 changed file with 12 additions and 4 deletions.
diff --git a/triton_viz/interpreter.py b/triton_viz/interpreter.py
@@ -343,12 +343,20 @@ def patch():
     old_create_dot = interpreter_builder.create_dot
     old_create_masked_store = interpreter_builder.create_masked_store
     GridExecutor.__call__ = _grid_executor_call
-    interpreter_builder.create_make_range = _create_make_range(interpreter_builder.create_make_range)
-    interpreter_builder.create_masked_load = _create_masked_load(interpreter_builder.create_masked_load)
-    interpreter_builder.create_expand_dims = _create_expand_dims(interpreter_builder.create_expand_dims)
+    interpreter_builder.create_make_range = _create_make_range(
+        interpreter_builder.create_make_range
+    )
+    interpreter_builder.create_masked_load = _create_masked_load(
+        interpreter_builder.create_masked_load
+    )
+    interpreter_builder.create_expand_dims = _create_expand_dims(
+        interpreter_builder.create_expand_dims
+    )
     interpreter_builder.binary_op = _create_binary_op(interpreter_builder.binary_op)
     interpreter_builder.create_dot = _create_dot(interpreter_builder.create_dot)
-    interpreter_builder.create_masked_store = _create_masked_store(interpreter_builder.create_masked_store)
+    interpreter_builder.create_masked_store = _create_masked_store(
+        interpreter_builder.create_masked_store
+    )
     try:
         yield
     finally: