Cleanup TE cuda graphs with the right api (#3459)

Signed-off-by: Gautham Kollu <gkollu@nvidia.com>
Signed-off-by: gautham-kollu <gkollu@nvidia.com>

Author: gautham-kollu

src/megatron/bridge/training/train.py

@@ -1567,11 +1567,7 @@ def _delete_cuda_graphs(cuda_graph_helper: TECudaGraphHelper):
 
     # Cleanup CUDA graphs object for partial Cuda-graphs (implemented in TransformerEngine)
     if cuda_graph_helper is not None:
-        for layers in cuda_graph_helper.callables_per_chunk:
-            for layer in layers:
-                for cuda_graph in layer.cuda_graphs:
-                    del cuda_graph
-                del layer.cuda_graphs
+        cuda_graph_helper.delete_cuda_graphs()
 
     # Run GC to collect the freshed object
     gc.collect()