added amp level option
This commit is contained in:
parent
35ca80683e
commit
e052a3bc92
|
@ -372,7 +372,8 @@ class Trainer(TrainerIO):
|
|||
if self.use_amp:
|
||||
for optimizer in self.optimizers:
|
||||
with amp.scale_loss(loss, optimizer) as scaled_loss:
|
||||
scaled_loss.backward()
|
||||
optimizer.backward(scaled_loss)
|
||||
# scaled_loss.backward()
|
||||
else:
|
||||
loss.backward()
|
||||
|
||||
|
|
2
setup.py
2
setup.py
|
@ -7,7 +7,7 @@ from setuptools import setup, find_packages
|
|||
# http://blog.ionelmc.ro/2014/05/25/python-packaging/
|
||||
setup(
|
||||
name="pytorch-lightning",
|
||||
version='0.1.dev1826',
|
||||
version='0.1.dev1827',
|
||||
description="The Keras for ML researchers using PyTorch",
|
||||
author="William Falcon",
|
||||
author_email="waf2107@columbia.edu",
|
||||
|
|
Loading…
Reference in New Issue