Last active
February 5, 2020 11:50
-
-
Save ihoromi4/dd2ed8d78f30ce8ccefe2f6bcb2871f1 to your computer and use it in GitHub Desktop.
pytorch - get nn.Module device method
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
from torch import nn | |
class NN(nn.Module): | |
def device(self): | |
return next(self.parameters()).device |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment