Merge pull request karpathy#236 from ymurenko/master

fix "cuda out of memory" when resuming training
Talos-1 · Apr 13, 2023 · 2457471 · 2457471
2 parents 553f949 + 4ac2e8c
commit 2457471
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/train.py b/train.py
@@ -190,6 +190,7 @@ def get_batch(split):
 optimizer = model.configure_optimizers(weight_decay, learning_rate, (beta1, beta2), device_type)
 if init_from == 'resume':
     optimizer.load_state_dict(checkpoint['optimizer'])
+checkpoint = None # free up memory
 
 # compile the model
 if compile: