fix imporerror for pytorch 2.9

Author: marksverdhei

examples/run_simple_mcore_train_loop.py

@@ -1,10 +1,10 @@
 import os
 import torch
 from torch.optim import Adam
-from torch.utils.data import DataLoader, Iterator
+from torch.utils.data import DataLoader
 from functools import partial
 from pathlib import Path
-from typing import Callable, Dict, Tuple
+from typing import Any, Callable, Dict, Tuple, Iterator
 
 from megatron.core import parallel_state
 from megatron.core import dist_checkpointing