changes from delete lazy [pr] (#8146)

* changes from delete lazy [pr]

* test tweak
This commit is contained in:
George Hotz
2024-12-10 11:06:17 -08:00
committed by GitHub
parent 286fec115e
commit aa3b094334
2 changed files with 11 additions and 10 deletions

View File

@@ -32,7 +32,7 @@ def helper_test(nm, gen, model, max_memory_allowed, max_kernels_allowed, all_jit
# TODO: jit should expose this correctly with graph
kernels_used = len(model.jit_cache) if hasattr(model, "jit_cache") else None
print(f"{nm}: used {mem_used/1e9:.2f} GB and {kernels_used} kernels in {min(tms)/1e6:.2f} ms")
assert mem_used/1e9 < max_memory_allowed, f"{nm} used more than {max_memory_allowed:.2f} GB"
assert mem_used/1e9 < max_memory_allowed, f"{nm} used more than {max_memory_allowed:.2f} GB - {mem_used/1e9:.2} GB used"
assert not kernels_used or kernels_used <= max_kernels_allowed, f"{nm} used more than {max_kernels_allowed} kernels"
if all_jitted:
assert kernels_used > 0 and kernels_used == GlobalCounters.kernel_count or (kernels_used <= GlobalCounters.kernel_count and getattr(Device[Device.DEFAULT], "graph", None)), f"only {kernels_used} out of {GlobalCounters.kernel_count} were jitted" # noqa: E501