From 4ac2e8ce3a47a38789dab355355ff822672a6d45 Mon Sep 17 00:00:00 2001
From: ymurenko <murenko.yuri@gmail.com>
Date: Wed, 5 Apr 2023 17:28:55 -0400
Subject: [PATCH] fix "cuda out of memory" when resuming training

---
 train.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/train.py b/train.py
index 30d0145..3c40524 100644
--- a/train.py
+++ b/train.py
@@ -189,6 +189,7 @@ scaler = torch.cuda.amp.GradScaler(enabled=(dtype == 'float16'))
 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: