Skip to content

Instantly share code, notes, and snippets.

@williamFalcon
Created July 29, 2019 17:51
Show Gist options
  • Save williamFalcon/aa7c67884b76ca7d0f2b0f50e2177276 to your computer and use it in GitHub Desktop.
Save williamFalcon/aa7c67884b76ca7d0f2b0f50e2177276 to your computer and use it in GitHub Desktop.
# change these lines
self.encoder = RNN(...)
self.decoder = RNN(...)
# to these
# now each RNN is based on a different gpu set
self.encoder = DataParallel(self.encoder, devices=[0, 1, 2, 3])
self.decoder = DataParallel(self.encoder, devices=[4, 5, 6, 7])
# in forward...
out = self.encoder(x.cuda(0))
# notice inputs on first gpu in device
sout = self.decoder(out.cuda(4)) # <--- the 4 here
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment