Thanks for the information!
You are currently pushing some tensors manually to a specific device:
if torch.cuda.is_available():
    indices = indices.to(device)
    probabilities = probabilities.to(device)

I’m not sure how device is defined, but I assume it’s the default GPU.
If you are using nn.DataParallel, you’ll get replicas on each GPU, so that the device won’t be static anymore inside your forward method.
Try to change each .to(device) call with a more generic one using the device of another parameter:
indices = indices.to(decoder_output.device)
probabilities = ...