wassname/pytorch-lightning
1
0
Fork 0
mirror of https://github.com/wassname/pytorch-lightning.git synced 2026-09-10 12:21:57 +08:00
Code Issues Packages Projects Releases Wiki Activity

fixed native amp + ddp (#1788)

Browse Source 
* fixed native amp + ddp

* fixed native amp + ddp
This commit is contained in:
William Falcon
2020-05-12 00:25:06 -04:00
committed by GitHub
parent 1df0d2dc97
commit 7b60d49432
1 changed files with 5 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pytorch_lightning/trainer/trainer.py
+5 -2
View File
@@ -510,8 +510,7 @@ class Trainer(
self.autocast_original_forward = None
self.use_native_amp = hasattr(torch.cuda, "amp") and hasattr(torch.cuda.amp, "autocast")
self.precision = precision
if self.use_native_amp and self.precision == 16:
self.scaler = torch.cuda.amp.GradScaler()
self.scaler = None
# TODO: remove for v0.8.0
self.amp_level = amp_level
@@ -858,6 +857,10 @@ class Trainer(
# set local properties on the model
self.copy_trainer_model_properties(ref_model)
# init amp. Must be done here instead of __init__ to allow ddp to work
if self.use_native_amp and self.precision == 16:
self.scaler = torch.cuda.amp.GradScaler()
# log hyper-parameters
if self.logger is not None:
# save exp to get started