diff --git a/gpytorch/priors/prior.py b/gpytorch/priors/prior.py index 4a1b454d6..c21ec252f 100644 --- a/gpytorch/priors/prior.py +++ b/gpytorch/priors/prior.py @@ -42,7 +42,11 @@ def __setattr__(self, name: str, value: Any) -> None: # If setting a BUFFERED_PREFIX attribute, update the base attribute instead. # Note: BUFFERED_PREFIX is just an indicator that this attribute belongs to a # TransformedDistribution, the value itself is not transformed. - if hasattr(self, name) and BUFFERED_PREFIX in name: + # + # Check BUFFERED_PREFIX before hasattr: torch.distributions lazy_property + # descriptors re-enter via setattr to cache (e.g. covariance_matrix), and + # hasattr() on those names would recurse infinitely. + if BUFFERED_PREFIX in name and hasattr(self, name): base_attr_name = name.replace(BUFFERED_PREFIX, "") # Convert to Tensor if needed tensor_value = torch.as_tensor(value) diff --git a/gpytorch/priors/torch_priors.py b/gpytorch/priors/torch_priors.py index 70b883160..b4f420e4b 100644 --- a/gpytorch/priors/torch_priors.py +++ b/gpytorch/priors/torch_priors.py @@ -152,6 +152,12 @@ def cpu(self): _del_attributes(module, MVN_LAZY_PROPERTIES) return module + def to(self, *args, **kwargs): + """Applies module-level to() call and resets all lazy properties""" + module = TModule.to(self, *args, **kwargs) + _del_attributes(module, MVN_LAZY_PROPERTIES) + return module + def expand(self, batch_shape): batch_shape = torch.Size(batch_shape) cov_shape = batch_shape + self.event_shape diff --git a/test/priors/test_multivariate_normal_prior.py b/test/priors/test_multivariate_normal_prior.py index e0612aac0..7390c675c 100644 --- a/test/priors/test_multivariate_normal_prior.py +++ b/test/priors/test_multivariate_normal_prior.py @@ -18,6 +18,19 @@ def test_multivariate_normal_prior_to_gpu(self): self.assertEqual(prior.scale_tril.device.type, "cuda") self.assertEqual(prior.precision_matrix.device.type, "cuda") + def test_multivariate_normal_prior_device_move_no_recursion(self): + # Accessing torch.distributions lazy properties after a device move must not + # RecursionError via Prior.__setattr__ / lazy_property caching. + prior = MultivariateNormalPrior(torch.tensor([0.0, 1.0]), covariance_matrix=torch.eye(2)) + _ = prior.covariance_matrix # cache lazy properties first + prior = prior.to(torch.device("cpu")) + self.assertEqual(prior.loc.device.type, "cpu") + self.assertEqual(prior.covariance_matrix.device.type, "cpu") + self.assertEqual(prior.scale_tril.device.type, "cpu") + self.assertEqual(prior.precision_matrix.device.type, "cpu") + prior = prior.cpu() + self.assertEqual(prior.covariance_matrix.device.type, "cpu") + def test_multivariate_normal_prior_validate_args(self): with self.assertRaises(ValueError): mean = torch.tensor([0.0, 1.0])