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

update (#7056)

This commit is contained in:
thomas chaton 2021-04-16 21:22:19 +01:00 committed by GitHub
parent 8bcd169767
commit 7b0b0d2844
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
pytorch_lightning/accelerators/accelerator.py
View File

@ -372,12 +372,7 @@ class Accelerator:
def to_device(self, batch: Any) -> Any:
"""Pushes the batch to the root device"""
# Todo (tchaton) Better fix
is_dict = isinstance(batch, dict)
if is_dict:
batch = [batch]
batch = self.batch_to_device(batch, self.root_device)
return batch[0] if is_dict else batch
return self.batch_to_device(batch, self.root_device)
@property
def amp_backend(self) -> Optional[LightningEnum]: