Rooba
/
lightning
mirror of https://github.com/Lightning-AI/lightning.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

added amp level option

This commit is contained in:
William Falcon 2019-05-16 15:52:00 -04:00
parent 35ca80683e
commit e052a3bc92
2 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
pytorch_lightning/models/trainer.py
View File

@ -372,7 +372,8 @@ class Trainer(TrainerIO):
if self.use_amp: if self.use_amp:
for optimizer in self.optimizers: for optimizer in self.optimizers:
with amp.scale_loss(loss, optimizer) as scaled_loss: with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward() optimizer.backward(scaled_loss)
# scaled_loss.backward()
else: else:
loss.backward() loss.backward()

2
setup.py
View File

@ -7,7 +7,7 @@ from setuptools import setup, find_packages
# http://blog.ionelmc.ro/2014/05/25/python-packaging/ # http://blog.ionelmc.ro/2014/05/25/python-packaging/
setup( setup(
name="pytorch-lightning", name="pytorch-lightning",
version='0.1.dev1826', version='0.1.dev1827',
description="The Keras for ML researchers using PyTorch", description="The Keras for ML researchers using PyTorch",
author="William Falcon", author="William Falcon",
author_email="waf2107@columbia.edu", author_email="waf2107@columbia.edu",