Remove extra gc cudas

mosaicml · Jul 21, 2024 · bc975cf · bc975cf
1 parent 6c071ee
commit bc975cf
Showing 1 changed file with 0 additions and 5 deletions.
diff --git a/llmfoundry/callbacks/hf_checkpointer.py b/llmfoundry/callbacks/hf_checkpointer.py
@@ -42,7 +42,6 @@
 from llmfoundry.utils.huggingface_hub_utils import \
     edit_files_for_hf_compatibility
 
-from llmfoundry.callbacks.scheduled_gc_callback import gc_cuda
 
 try:
     import transformer_engine.pytorch as te
@@ -395,8 +394,6 @@ def _save_checkpoint(self, state: State, logger: Logger):
 
         log.debug(f'memory before first gc {psutil.virtual_memory()}')
 
-        gc_cuda()
-
         log.debug(f'memory {psutil.virtual_memory()}')
 
         if state.is_model_ddp:
@@ -492,8 +489,6 @@ def dtensor_to_tensor_hook(
         log.debug(f'memory after gather {psutil.virtual_memory()}')
 
 
-        gc_cuda()
-
 
         # you can have the percentage of used RAM
         log.debug(f'memory after gather after gc {psutil.virtual_memory()}')