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

Device (#1790)

Browse Source 
* added self.device

* added docs
This commit is contained in:
William Falcon
2020-05-12 00:09:48 -04:00
committed by GitHub
parent de1fdd8d3b
commit 4b30ef6480
5 changed files with 13 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/source/multi_gpu.rst
+3 -1
View File
@@ -46,7 +46,9 @@ This will make your code scale to any arbitrary number of GPUs or TPUs with Ligh
# with lightning
def forward(self, x):
z = torch.Tensor(2, 3)
z = z.type_as(x)
z = z.type_as(x, device=self.device)
Every LightningModule knows what device it is on. You can access that reference via `self.device`.
Remove samplers
^^^^^^^^^^^^^^^