osmarks/nanogpt-experiments
1
0
Fork 0
mirror of https://github.com/osmarks/nanogpt-experiments.git synced 2024-12-18 06:00:29 +00:00
Code Issues Releases Wiki Activity

very slight refactor, bit cleaner

Browse Source
This commit is contained in:
Andrej Karpathy 2023-02-04 19:34:24 +00:00
parent dc149891b6
commit e108ffb973
1 changed files with 2 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
train.py
View File

@ -112,12 +112,10 @@ def get_batch(split):
ix = torch.randint(len(data) - block_size, (batch_size,))
x = torch.stack([torch.from_numpy((data[i:i+block_size]).astype(np.int64)) for i in ix])
y = torch.stack([torch.from_numpy((data[i+1:i+1+block_size]).astype(np.int64)) for i in ix])
# pin arrays x,y, which allows us to move them to GPU asynchronously (non_blocking=True)
if "cuda" in device:
# GPU training
if device_type == 'cuda':
# pin arrays x,y, which allows us to move them to GPU asynchronously (non_blocking=True)
x, y = x.pin_memory().to(device, non_blocking=True), y.pin_memory().to(device, non_blocking=True)
else:
# CPU or MPS training
x, y = x.to(device), y.to(device)
return x, y