I found a way around it:
def tr2Dist_new(r):
    nb,_,n = r.shape
    d = torch.zeros((nb,n,n),dtype=r.dtype,device=r.device)
    tmp = torch.sum(r ** 2, dim=1).unsqueeze(1) + torch.sum(r ** 2, dim=1).unsqueeze(2) - 2 * r.transpose(1,2) @ r
    m = torch.ones_like(d,dtype=torch.bool,device=r.device)
    m[..., torch.arange(m.shape[-1]), torch.arange(m.shape[-1])] = False
    m2 = tmp > 0
    m3 = m*m2
    tmp2 = tmp[m3]
    d[m3] = torch.sqrt(tmp2)
    return d

But it definitely makes the code more complicated and a bit slower I guess.
I found a similar problem further on in my loss function, where I have to loop over the individual samples in a batch in order to get around it as far as I can see:
        if safeversion:
            result = 0
            for i in range(nb):
                inputi = input[i,...]
                targeti = target[i,...]
                maski = mask[i,...]
                result += torch.sum(torch.norm(inputi[maski] - targeti[maski]) ** 2 / torch.norm(targeti[maski]) ** 2)
        else:
            result = torch.sum(torch.norm(input * mask - target * mask,dim=(1,2)) ** 2 / torch.norm(target * mask,dim=(1,2)) ** 2)

It works, but it is convoluted and quite a bit slower. I’m really puzzled that other people don’t seem to have trouble with this?
Thanks for the help albanD