Supress warning for the x_test equal to x_train

pycroscopy · Jun 18, 2021 · a25aac2 · a25aac2
1 parent ae3623d
commit a25aac2
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/atomai/models/dklgp/dklgpr.py b/atomai/models/dklgp/dklgpr.py
@@ -7,6 +7,7 @@
 Created by Maxim Ziatdinov (email: maxim.ziatdinov@ai4microscopy.com)
 """
 
+import warnings
 from typing import Tuple, Type, Union
 
 import gpytorch
@@ -74,8 +75,11 @@ def _compute_posterior(self, X: torch.Tensor) -> gpytorch.distributions.Multivar
         """
         self.gp_model.eval()
         self.likelihood.eval()
-        with torch.no_grad(), gpytorch.settings.use_toeplitz(False), gpytorch.settings.fast_pred_var():
-            posterior = self.gp_model(X.to(self.device))
+        wrn = gpytorch.models.exact_gp.GPInputWarning
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=wrn)
+            with torch.no_grad(), gpytorch.settings.use_toeplitz(False), gpytorch.settings.fast_pred_var():
+                posterior = self.gp_model(X.to(self.device))
         return posterior
 
     def sample_from_posterior(self, X, num_samples: int = 1000) -> np.ndarray: