Commit aba7eaf1 authored by Hsing-Yin Chang's avatar Hsing-Yin Chang
Browse files

fix gpu_task_wrapper

parent 8a982a04
......@@ -7,20 +7,20 @@ from pygion import task, Region, Partition, Tunable, R
from spinifel import SpinifelContexts
def gpu_task_wrapper(thunk):
@wraps(thunk)
def wrapper(*args, **kwargs):
context = SpinifelContexts()
context.ctx.push()
if context.ctx is not None:
context.ctx.push()
try:
return thunk(*args, **kwargs)
finally:
context.ctx.pop()
if context.ctx is not None:
context.ctx.pop()
return wrapper
@task(privileges=[R])
@nvtx.annotate("legion/utils.py", is_prefix=True)
def print_region(region):
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment