[train] alias train.Checkpoint to new Checkpoint (ray-project#38864)

* [train] alias train.Checkpoint

Signed-off-by: Matthew Deng <matt@anyscale.com>

python/ray/train/__init__.py

@@ -19,7 +19,13 @@
 from ray.train.context import get_context
 from ray.train.trainer import TrainingIterator
 
-from ray.air import Checkpoint
+from ray.train._internal.storage import _use_storage_context
+
+if _use_storage_context():
+    from ray.train._checkpoint import Checkpoint
+else:
+    from ray.air import Checkpoint
+
 from ray.air.config import CheckpointConfig, FailureConfig, RunConfig, ScalingConfig
 from ray.air.result import Result