From c636f5bd40187bfcde184979102722fc6fd80905 Mon Sep 17 00:00:00 2001 From: Amog Kamsetty Date: Mon, 21 Sep 2020 13:10:10 -0700 Subject: [PATCH] [Ray SGD] FP16 Hotfix (#10931) --- python/ray/util/sgd/torch/training_operator.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/ray/util/sgd/torch/training_operator.py b/python/ray/util/sgd/torch/training_operator.py index 85ba958c0..aad6db064 100644 --- a/python/ray/util/sgd/torch/training_operator.py +++ b/python/ray/util/sgd/torch/training_operator.py @@ -257,8 +257,8 @@ class TrainingOperator: if self.use_fp16 and amp: logger.debug("Setting up Apex.") - self._models, self._optimizers = amp.initialize( - self._models, self._optimizers, **self._apex_args) + self._original_models, self._optimizers = amp.initialize( + self._original_models, self._optimizers, **self._apex_args) self._amp = amp if self._wrap_ddp: