[tune] fix pbt ptl example (ray-project#11573)

* [tune] fix pbt ptl example * wider smoke test
krfricke · Oct 23, 2020 · 8ee4f7e · 8ee4f7e
1 parent 7a0184e
commit 8ee4f7e
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/python/ray/tune/examples/mnist_pytorch_lightning.py b/python/ray/tune/examples/mnist_pytorch_lightning.py
@@ -248,7 +248,7 @@ def tune_mnist_pbt(num_samples=10, num_epochs=10, gpus_per_trial=0):
  mode="min",
  perturbation_interval=4,
  hyperparam_mutations={
- "lr": lambda: tune.loguniform(1e-4, 1e-1).func(None),
+ "lr": tune.loguniform(1e-4, 1e-1),
  "batch_size": [32, 64, 128]
  })
 
@@ -285,8 +285,8 @@ def tune_mnist_pbt(num_samples=10, num_epochs=10, gpus_per_trial=0):
  args, _ = parser.parse_known_args()
 
  if args.smoke_test:
- tune_mnist_asha(num_samples=1, num_epochs=1, gpus_per_trial=0)
- tune_mnist_pbt(num_samples=1, num_epochs=1, gpus_per_trial=0)
+ tune_mnist_asha(num_samples=1, num_epochs=6, gpus_per_trial=0)
+ tune_mnist_pbt(num_samples=1, num_epochs=6, gpus_per_trial=0)
  else:
  # ASHA scheduler
  tune_mnist_asha(num_samples=10, num_epochs=10, gpus_per_trial=0)