[SGD] Fix Docstring for as_trainable (#13173)

2026-06-27 20:06:31 +08:00 · 2021-01-04 17:21:24 -08:00
parent 15e86581bd
commit e181515dff
1 changed files with 4 additions and 2 deletions
@@ -594,7 +594,8 @@ class TorchTrainer:

            TorchTrainable = TorchTrainer.as_trainable(
                training_operator_cls=MyTrainingOperator,
-                num_gpus=2,
+                num_workers=2,
+                use_gpu=True,
                override_tune_step=step
            )
            analysis = tune.run(
@@ -695,7 +696,8 @@ class BaseTorchTrainable(Trainable):
        # TorchTrainable is subclass of BaseTorchTrainable.
        TorchTrainable = TorchTrainer.as_trainable(
            training_operator_cls=MyTrainingOperator,
-            num_gpus=2,
+            num_workers=2,
+            use_gpu=True,
            override_tune_step=custom_step
        )