import cupy as cp
from contextlib import contextmanager

@contextmanager
def sync_time(name):
    start = cp.cuda.Event()
    end = cp.cuda.Event()
    start.record()
    start.synchronize()
    yield
    end.record()
    end.synchronize()
    t = cp.cuda.get_elapsed_time(start,end)
    print("{} : {} ms".format(name,t))

x = cp.random.normal(size=((400, 32, 28, 28))).astype(cp.float32)

with sync_time("cupy"):
    for i in range(1000):
        x.max()

x = cp.asnumpy(x) #move to cpu

with sync_time("numpy"):
    for i in range(1000):
        x.max()