[release/2.5] ModuleTracker: Add explicit garbage collection

When running an FSDP model with FlopCounterMode, we are experiencing a memory
leak. It is coming from ModuleTracker class. Even though
ModuleTracker class is keeping weakrefrences of the operators, the
tensors/operators are not being freed after the backward pass. To force
free these tensors/operators after backwardpass, I explicitly added
garbage collection in the post forward hook.

(cherry picked from commit 63dc40d)

torch/utils/module_tracker.py

@@ -10,7 +10,7 @@
     register_module_forward_pre_hook,
 )
 from torch.utils._pytree import tree_flatten
-
+import gc
 
 logger = logging.getLogger(__name__)
 
@@ -136,6 +136,7 @@ def _fw_post_hook(self, mod, input, output):
         tensors = [a for a in args if isinstance(a, torch.Tensor) and a.requires_grad]
         if tensors:
             register_multi_grad_hook(tensors, self._get_append_fn(name, True))
+            gc.collect()
 
     def __enter__(self):
         self._fw_pre_handle = register_module_forward_pre_hook(self._fw_pre_hook)