[release] minor fix to pytorch_pbt_failure test when using gpu.
diff --git a/python/ray/train/examples/pytorch/tune_cifar_torch_pbt_example.py b/python/ray/train/examples/pytorch/tune_cifar_torch_pbt_example.py
index 90846eb84824e..196051971129f 100644
--- a/python/ray/train/examples/pytorch/tune_cifar_torch_pbt_example.py
+++ b/python/ray/train/examples/pytorch/tune_cifar_torch_pbt_example.py
@@ -70,6 +70,10 @@ def train_func(config):
 
     model = resnet18()
 
+    # Note that `prepare_model` needs to be called before setting optimizer.
+    if not session.get_checkpoint():  # fresh start
+        model = train.torch.prepare_model(model)
+
     # Create optimizer.
     optimizer_config = {
         "lr": config.get("lr"),
@@ -84,6 +88,7 @@ def train_func(config):
         # Load in model
         model_state = checkpoint_dict["model"]
         model.load_state_dict(model_state)
+        model = train.torch.prepare_model(model)
 
         # Load in optimizer
         optimizer_state = checkpoint_dict["optimizer_state_dict"]
@@ -97,8 +102,6 @@ def train_func(config):
         checkpoint_epoch = checkpoint_dict["epoch"]
         starting_epoch = checkpoint_epoch + 1
 
-    model = train.torch.prepare_model(model)
-
     # Load in training and validation data.
     transform_train = transforms.Compose(
         [

manually tested on https://console.anyscale-staging.com/o/anyscale-internal/projects/prj_qC3ZfndQWYYjx2cz8KWGNUL4/clusters/ses_wttwryzcyl6ahytbfz23vwm5fa?command-history-section=head_start_up_log

Signed-off-by: xwjiang2010 <xwjiang2010@gmail.com>

<!-- Thank you for your contribution! Please review https://github.com/ray-project/ray/blob/master/CONTRIBUTING.rst before opening a pull request. -->

<!-- Please add a reviewer to the assignee section when you create a PR. If you don't have the access to it, we will shortly find a reviewer and assign them to your PR. -->

## Why are these changes needed?

<!-- Please give a short summary of the change and the problem this solves. -->

## Related issue number

<!-- For example: "Closes #1234" -->

## Checks

- [ ] I've signed off every commit(by using the -s flag, i.e., `git commit -s`) in this PR.
- [ ] I've run `scripts/format.sh` to lint the changes in this PR.
- [ ] I've included any doc changes needed for https://docs.ray.io/en/master/.
- [ ] I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
- Testing Strategy
   - [ ] Unit tests
   - [ ] Release tests
   - [ ] This PR is not tested :(
