File tree Expand file tree Collapse file tree 1 file changed +2
-1
lines changed
src/pytorch_lightning/loops/optimization Expand file tree Collapse file tree 1 file changed +2
-1
lines changed Original file line number Diff line number Diff line change 31
31
_extract_hiddens ,
32
32
)
33
33
from pytorch_lightning .plugins import ApexMixedPrecisionPlugin
34
+ from pytorch_lightning .plugins .precision .native_amp import MixedPrecisionPlugin
34
35
from pytorch_lightning .trainer .progress import OptimizationProgress
35
36
from pytorch_lightning .utilities .exceptions import MisconfigurationException
36
37
from pytorch_lightning .utilities .rank_zero import rank_zero_deprecation
@@ -366,7 +367,7 @@ def _optimizer_step(
366
367
" `using_native_amp` argument. Removing this argument will avoid this message, you can expect it to"
367
368
" return True."
368
369
)
369
- kwargs ["using_native_amp" ] = self .trainer .amp_backend == "native"
370
+ kwargs ["using_native_amp" ] = isinstance ( self .trainer .precision_plugin , MixedPrecisionPlugin )
370
371
self .trainer ._call_lightning_module_hook (
371
372
"optimizer_step" ,
372
373
self .trainer .current_epoch ,
You can’t perform that action at this time.
0 commit comments