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

Mem crash (#299)

Browse Source 
* fixes memory crash

* fixes memory crash
This commit is contained in:
William Falcon
2019-10-04 15:53:44 -04:00
committed by GitHub
parent 36f0b5bbd0
commit 73a7cf3c99
2 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pytorch_lightning/root_module/memory.py
+1 -1
View File
@@ -35,7 +35,7 @@ class ModelSummary(object):
out_sizes = []
input_ = self.model.example_input_array
if self.model.use_ddp:
if self.model.use_ddp or self.model.use_dp:
input_ = input_.cuda(0)
if self.model.trainer.use_amp: