Update train.py

train.py

@@ -192,10 +192,9 @@ def main(args):
     # transfer the train data loader into train data iterator
     train_iter = iter(train_dl)
 
-    # close automatic garbage collection
-    gc.disable()
-
     with initialize_llm_profile(profiling=args.profiling, start_time=current_time) as prof:
+        # close automatic garbage collection
+        gc.disable()
         # start iterating the train data and begin training
         for batch_count in range(train_state.batch_count, total_steps):
             empty_cache_and_diag(batch_count, interval=gpc.config.data.empty_cache_and_diag_interval)