Use postprocess func for grad kernel
#!/usr/bin/env python3
-from .rbf_kernel import RBFKernel
+from .rbf_kernel import RBFKernel, postprocess_rbf
import torch
from ..lazy.kronecker_product_lazy_tensor import KroneckerProductLazyTensor
@@ -67,8 +67,14 @@ class RBFKernelGrad(RBFKernel):
outer = torch.transpose(outer, -1, -2).contiguous()
# 1) Kernel block
- diff = self.covar_dist(x1_, x2_, square_dist=True, **params)
- K_11 = diff.div_(-2).exp_()
+ diff = self.covar_dist(
+ x1_,
+ x2_,
+ square_dist=True,
+ dist_postprocess_func=postprocess_rbf,
+ **params
+ )
+ K_11 = diff
K[..., :n1, :n2] = K_11
# 2) First gradient block
