Ensure that GC remains disabled in GC callback (#534)

Summary:
Pull Request resolved: #534

Ensure that the GC remains disabled after every training batch since it may have been re-enabled elsewhere.

Reviewed By: anshulverma

Differential Revision: D49100568

fbshipit-source-id: f3dae95745778face82aa26af6fe06dd09dab50f

torchtnt/framework/callbacks/garbage_collector.py

@@ -56,6 +56,9 @@ def on_train_step_end(self, state: State, unit: TTrainUnit) -> None:
         if total_num_steps_completed % self._step_interval == 0:
             gc.collect()
 
+        # Ensure that GC is disabled, in case GC was reenabled elsewhere
+        gc.disable()
+
     def on_train_end(self, state: State, unit: TTrainUnit) -> None:
         gc.enable()