osmarks/nanogpt-experiments
1
0
Fork 0
mirror of https://github.com/osmarks/nanogpt-experiments.git synced 2024-11-13 05:19:58 +00:00
Code Issues Releases Wiki Activity

pull out dtype up top

Browse Source
This commit is contained in:
Andrej Karpathy 2022-12-29 05:32:55 +00:00
parent e7bac659f5
commit fa57d464d7
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
bench.py
View File

@ -14,6 +14,7 @@ torch.manual_seed(1337)
batch_size = 8
block_size = 1024
dtype = torch.float16
# data loading init
real_data = True
@ -65,7 +66,7 @@ if profile:
for k in range(num_steps):
X, Y = get_batch('train')
with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
with torch.autocast(device_type='cuda', dtype=dtype):
logits, loss = model(X, Y)
optimizer.zero_grad(set_to_none=True)
loss.backward()
@ -83,7 +84,7 @@ else:
t0 = time.time()
for k in range(num_steps):
X, Y = get_batch('train')
with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
with torch.autocast(device_type='cuda', dtype=dtype):
logits, loss = model(X, Y)
optimizer.zero_grad(set_to_none=True)
loss.backward()