I'm quite new to using torch but maybe you can try setting your device to something like this it first checks for "mps" if its not available it checks GPU "cuda" and if that's not available finally the CPU.

        self.device = (
            "mps"
            if torch.backends.mps.is_available() and torch.backends.mps.is_built()
            else ("cuda" if torch.cuda.is_available() else "cpu")
        )

since you have 2 gpu's cuda:0 and cuda:1 wrap your model with DataParallel and send it to the devices defined above. Hope this helps but as i said i am very new to using torch and i have no idea what's being pulled in with those templates or auto mode ect.

        model = torch.nn.DataParallel(model)
        model.to(self.device)

** Note: DataParallel is being considered for deprecation so if you can figure it out DistributedDataParallel that would be a better solution but a bit more complex.
also make sure there**