Skip to content

Commit 3630756

Browse files
apaszkeGoogle-ML-Automation
authored andcommitted
[Mosaic GPU] Use events as the default profiling method
JAX still supports old CUDA versions such as 12.0, where CUPTI leaks memory. PiperOrigin-RevId: 705459909
1 parent dda6b88 commit 3630756

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

jax/experimental/mosaic/gpu/profiler.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -114,7 +114,7 @@ def wrapper(*args, **kwargs):
114114
return wrapper
115115

116116

117-
def measure(f: Callable, *, mode: str = "cupti", aggregate: bool = True
117+
def measure(f: Callable, *, mode: str = "events", aggregate: bool = True
118118
) -> Callable:
119119
"""Sets up a function ``f`` for profiling on GPU.
120120

0 commit comments

Comments
 (0)