-
Notifications
You must be signed in to change notification settings - Fork 24.7k
Description
🐛 Describe the bug
The torch.gather and Tensor.gather documentation states that input and index must have the same number of dimensions. However, no corresponding validation is added.Validation needs to be added for both CPU and GPU.
cpu:
import torch
input_data = torch.tensor([[1, 2], [3, 4], [5, 6]])
indices = torch.tensor([])
output = torch.gather(input_data, dim=0, index=indices)
print(output)
gpu:
import time
import math
import numpy as np
import torch
input_data = torch.tensor([[1, 2], [3, 4], [5, 6]])
input_data = input_data.cuda()
indices = torch.tensor([])
indices = indices.cuda()
output = torch.gather(input_data, dim=0, index=indices)
output = output.cuda()
npu:
import time
import math
import numpy as np
import torch
import torch_npu
input_data = torch.tensor([[1, 2], [3, 4], [5, 6]])
input_data = input_data.npu()
indices = torch.tensor([])
indices = indices.npu()
output = torch.gather(input_data, dim=0, index=indices)
output = output.npu()
Versions
RuntimeError: call aclnnGather failed, detail:EZ1001: [PID: 64283] 2025-05-21-07:32:10.823.227 Index tensor must have the same number of dimensions as input tensor.
cc @malfet