Transferring Data Augmentation Modules to Device #19395
Unanswered
EthanMarx
asked this question in
DDP / multi-GPU / multi-node
Replies: 0 comments
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
-
Hello! In my
DataModule
I setup a couple oftorch.nn.Modules
that do on-gpu data augmentations. I am wondering how to ensure that these Modules get transferred to the appropriate device when doing multi gpu training?My setup looks something like the pseudo code below
appreciate the help
Beta Was this translation helpful? Give feedback.
All reactions