Это лучший способ для вычисления ядра в последней версии PyTorch (преобразование TensorFlow в PyTorch)?

Question

Я пытаюсь преобразовать реализацию MMD-VAE из TensorFlow в PyTorch.У меня большая часть модели построена просто отлично, но я просто хочу убедиться, что я правильно конвертирую следующие функции (все работает, но я не получаю ожидаемых результатов, поэтому подумал, что, возможно, неправильно вычисляю ядротак как я не так силен в TensorFlow)

def compute_kernel(x, y):
    x_size = x.size(0)
    y_size = y.size(0)
    dim = x.size(1)
    x = x.unsqueeze(1)
    y = y.unsqueeze(0)
    tiled_x = x.expand(x_size, y_size, dim)
    tiled_y = y.expand(x_size, y_size, dim)
    kernel_input = (tiled_x - tiled_y).pow(2).mean(2)/float(dim)
    return torch.exp(-kernel_input)

Спасибо за вашу помощь !!