diff --git a/src/transformers/models/perceiver/modeling_perceiver.py b/src/transformers/models/perceiver/modeling_perceiver.py index 40d8f8fb9970..00ec22ba3fc3 100755 --- a/src/transformers/models/perceiver/modeling_perceiver.py +++ b/src/transformers/models/perceiver/modeling_perceiver.py @@ -1978,7 +1978,7 @@ def decoder_query(self, inputs, modality_sizes=None, inputs_without_pos=None, su # to get the indices for the unflattened array # unravel_index returns a tuple (x_idx, y_idx, ...) # stack to get the [n, d] tensor of coordinates - indices = torch.unravel_index(subsampled_points, self.output_index_dims) + indices = torch.unravel_index(subsampled_points.to(inputs.device), self.output_index_dims) pos = torch.stack(indices, dim=1) batch_size = inputs.shape[0] # Map these coordinates to [-1, 1]