AssertionError in models.py Group module line assert idx.size(1) == self.num_group dimension mismatch error

[SouthernBadWood]

model = create_fn(
Extracting train features for class bagel: 0%| | 0/244 [00:02<?, ?it/s]
Traceback (most recent call last):
File "/home/data/cwl/M3DM/main.py", line 123, in
run_3d_ads(args)
File "/home/data/cwl/M3DM/main.py", line 19, in run_3d_ads
model.fit(cls)
File "/home/data/cwl/M3DM/m3dm_runner.py", line 47, in fit
method.add_sample_to_mem_bank(sample)
File "/home/data/cwl/M3DM/feature_extractors/multiple_features.py", line 700, in add_sample_to_mem_bank
rgb_feature_maps, xyz_feature_maps, center, neighbor_idx, center_idx, interpolated_pc = self(sample[0],unorganized_pc_no_zeros.contiguous())
File "/home/data/cwl/M3DM/feature_extractors/features.py", line 102, in call
rgb_feature_maps, xyz_feature_maps, center, ori_idx, center_idx = self.deep_feature_extractor(rgb, xyz)
File "/home/data/miniconda3/envs/pcd/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
return forward_call(*input, **kwargs)
File "/home/data/cwl/M3DM/models/models.py", line 54, in forward
xyz_features, center, ori_idx, center_idx = self.xyz_backbone(xyz)
File "/home/data/miniconda3/envs/pcd/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
return forward_call(*input, **kwargs)
File "/home/data/cwl/M3DM/models/models.py", line 347, in forward
neighborhood, center, ori_idx, center_idx = self.group_divider(pts)
File "/home/data/miniconda3/envs/pcd/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
return forward_call(*input, **kwargs)
File "/home/data/cwl/M3DM/models/models.py", line 86, in forward
assert idx.size(1) == self.num_group
AssertionError

ther was a dimension mismatch error in assert idx.size(1) == self.num_group # 1024 vs 128 Has anyone encountered similar dimension mismatch issues，I would greatly appreciate it if you could provide any assistance.

[SouthernBadWood]

I have successfully resolved this issue. The problem was caused by the version of KNN. If you encounter the same issue, you can update your KNN by running the following command:
pip install --upgrade https://github.com/unlimblue/KNN_CUDA/releases/download/0.2/KNN_CUDA-0.2-py3-none-any.whl