RuntimeErrors with "torch_two_sample.statistics_diff.SmoothKNNStatistic"

[adrienchaton]

Hello !
Thank you very much for sharing these functions to the PyTorch community,

I am successfully using the "torch_two_sample.statistics_diff.MMDStatistic" both at training/backprop and evaluation, running on GPU.

I am trying to have an alternative criterion, using instead the SmoothKNN that I set as the MMD (with the additional True boolean for cuda and the k parameter).

The code is absolutely identical to when using the MMD for the criterion but with the SmoothKNN function it gives a "RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation"

Does anyone have a fix to that please ?

Thanks in advance !

[rekcahpassyla]

May be a bit late for you but I fixed it by changing statistics_diff.py in the SmoothKNNStatistic.__call__ function as follows (showing the lines that were commented out and replaced)

        if margs_.is_cuda:
            #margs.masked_scatter_(indices, margs_.view(-1))
            margs.masked_scatter_(indices, margs_.detach().view(-1))
        else:
            #margs.masked_scatter_(indices_cpu, margs_.view(-1))
            margs.masked_scatter_(indices_cpu, margs_.detach().view(-1))