github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-04-29 03:00:14 -04:00
Code Issues Packages Projects Releases Wiki Activity

err, okay, a bit more off

Browse Source
This commit is contained in:
George Hotz
2022-06-24 22:44:57 -07:00
parent bdde95f16e
commit d748353ce5
1 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
examples/benchmark_train_efficientnet.py
View File

@@ -33,11 +33,12 @@ if __name__ == "__main__":
loss = out.logsoftmax().mul(y_train).mean()
optimizer.zero_grad()
loss.backward()
mt = time.monotonic()
loss = loss.cpu().data[0]
et = time.monotonic()
loss = loss.cpu().data[0]
info = pynvml.nvmlDeviceGetMemoryInfo(handle)
print(f"{(et-st)*1000.0:7.2f} ms, {loss:7.2f} loss, {tensors_allocated():4d} tensors, {info.used/1e9:.2f} GB used")
print(f"{(et-st)*1000.0:7.2f} ms run, {(mt-st)*1000.0:7.2f} ms build, {loss:7.2f} loss, {tensors_allocated():4d} tensors, {info.used/1e9:.2f} GB used")