When I just use UmbrellaSurfaceConstructor in front of other networks, the training of the whole network becomes very slow

[ybc-ybc]

When I just use UmbrellaSurfaceConstructor in front of other networks, the training of the whole network becomes very slow, especially the loss stage, do you know the reason?
The data size is the same as Repsurf, but the speed is much slower.
Remove UmbrellaSurfaceConstructor, the speed will recover quickly.

[hancyran]

Can you provide some code snippets for this problem? Do you use the cls's UmbrellaSurfaceConstructor or the seg's? I may consider a reason that if you do not use cuda-ops for kNN (like the given lib pointops), the speed will be quite slow.