Error when inputting pruned sparse tensor to MinkowskiUNet

[jchibane]

Hello,

when I input a sparse tensor that was previously pruned using ME.MinkowskiPruning() into a minkowski U-Net (like MinkUNet34C) I get an error:

AssertionError: SparseTensors must have the same coordinate_map_key.coordinate map key:[8, 8, 8] != coordinate map key:[8, 8, 8]:pruned-4cBWD

at the first ME.cat operation in the U-Net decoder. Before the first transpose convolution, the annotation "pruned-XXXX" is attached to the coordinate map key but is missing thereafter. It seems this throws the error. One workaround would be to do the masking coordinates and features as pytorch tensors manually and then recreating a sparse tensor. But I'm wondering if this will preserve backward gradient flow, will be efficient or there is a better way. Best, Julian

ME version 0.5.4

[jchibane]

Doing the naive thing, i.e. the masking on coordinates and features as regular pytorch tensor and recreating a sparse tensor seems to work and preserve gradient flow. So as a workaround that seems valid. Best, Julian