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

Merge pull request #236 from ymurenko/master

Browse Source 
fix "cuda out of memory" when resuming training
This commit is contained in:
Andrej 2023-04-12 22:09:42 -07:00 committed by GitHub
commit 2457471c9c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
train.py
View File

@ -190,6 +190,7 @@ scaler = torch.cuda.amp.GradScaler(enabled=(dtype == 'float16'))
optimizer = model.configure_optimizers(weight_decay, learning_rate, (beta1, beta2), device_type) optimizer = model.configure_optimizers(weight_decay, learning_rate, (beta1, beta2), device_type)
if init_from == 'resume': if init_from == 'resume':
optimizer.load_state_dict(checkpoint['optimizer']) optimizer.load_state_dict(checkpoint['optimizer'])
checkpoint = None # free up memory
# compile the model # compile the model
if compile: if compile: