[train] fix tensorflow example by using ScalingConfig (#46565)

Signed-off-by: MJ <[email protected]>

python/ray/train/examples/tf/tensorflow_autoencoder_example.py

@@ -14,7 +14,7 @@
 from ray.air.integrations.keras import ReportCheckpointCallback
 from ray.data.datasource import SimpleTensorFlowDatasource
 from ray.data.extensions import TensorArray
-from ray.train import Result
+from ray.train import Result, ScalingConfig
 from ray.train.tensorflow import TensorflowTrainer, prepare_dataset_shard
 
 
@@ -121,7 +121,7 @@ def train_tensorflow_mnist(
 ) -> Result:
     train_dataset = get_dataset(split_type="train")
     config = {"lr": 1e-3, "batch_size": 64, "epochs": epochs}
-    scaling_config = dict(num_workers=num_workers, use_gpu=use_gpu)
+    scaling_config = ScalingConfig(num_workers=num_workers, use_gpu=use_gpu)
     trainer = TensorflowTrainer(
         train_loop_per_worker=train_func,
         train_loop_config=config,

python/ray/train/examples/tf/tune_tensorflow_autoencoder_example.py

@@ -2,6 +2,7 @@
 
 import ray
 from ray import tune
+from ray.train import ScalingConfig
 from ray.train.examples.tf.tensorflow_mnist_example import train_func
 from ray.train.tensorflow import TensorflowTrainer
 from ray.tune.tune_config import TuneConfig
@@ -11,7 +12,7 @@
 def tune_tensorflow_mnist(
     num_workers: int = 2, num_samples: int = 2, use_gpu: bool = False
 ):
-    scaling_config = dict(num_workers=num_workers, use_gpu=use_gpu)
+    scaling_config = ScalingConfig(num_workers=num_workers, use_gpu=use_gpu)
     trainer = TensorflowTrainer(
         train_loop_per_worker=train_func,
         scaling_config=scaling_config,