Merge pull request #2378 from j-wilson/patch_matern

Minor patch to Matern covariances

Author: Balandat

gpytorch/functions/matern_covariance.py

@@ -13,7 +13,7 @@ def forward(ctx, x1, x2, lengthscale, nu, dist_func):
         # Subtract mean for numerical stability. Won't affect computations
         # because covariance matrix is stationary.
         needs_grad = any(ctx.needs_input_grad)
-        mean = x1.reshape(-1, x1.size(-1)).mean(0)[(None,) * (x1.dim() - 1)]
+        mean = x1.mean(dim=-2, keepdim=True)
         x1_ = (x1 - mean).div(lengthscale)
         x2_ = (x2 - mean).div(lengthscale)
         scaled_unitless_dist = dist_func(x1_, x2_).mul_(math.sqrt(2 * nu))

gpytorch/kernels/matern_kernel.py

@@ -92,7 +92,7 @@ def forward(self, x1, x2, diag=False, **params):
             or params.get("last_dim_is_batch", False)
             or trace_mode.on()
         ):
-            mean = x1.reshape(-1, x1.size(-1)).mean(0)[(None,) * (x1.dim() - 1)]
+            mean = x1.mean(dim=-2, keepdim=True)
 
             x1_ = (x1 - mean).div(self.lengthscale)
             x2_ = (x2 - mean).div(self.lengthscale)