I am using robustness library and would like to know how one should use multiple GPUs with it, specifically I want to train the robust model and I have the model as torch vision V2 and I have modified the forward function to work with the library but still I get errors related to input and model not being on the same device. On fixing it, it stills says that for the loss function the output and the target are on the Cuda:1 and Cuda:0. So is there something I am missing?
I am using robustness library and would like to know how one should use multiple GPUs with it, specifically I want to train the robust model and I have the model as torch vision V2 and I have modified the forward function to work with the library but still I get errors related to input and model not being on the same device. On fixing it, it stills says that for the loss function the output and the target are on the Cuda:1 and Cuda:0. So is there something I am missing?