Make some tensors buffers in pathwise sampling code (#1821)

Summary:
Fixes an issue when moving a SamplePath from GPU to GPU.

I wasn't super thorough to go through all of the code and every eventuality, but this at least makes the issue reported in #1820 go away.

Pull Request resolved: #1821

Reviewed By: esantorella

Differential Revision: D45695243

Pulled By: Balandat

fbshipit-source-id: 561131330badc545abae88cfe05673cc518c8296

botorch/sampling/pathwise/features/maps.py

@@ -109,6 +109,8 @@ def __init__(
         """
         super().__init__()
         self.kernel = kernel
+        self.register_buffer("weight", weight)
+        self.register_buffer("bias", bias)
         self.weight = weight
         self.bias = bias
         self.input_transform = input_transform

botorch/sampling/pathwise/paths.py

@@ -174,6 +174,8 @@ def __init__(
         """
         super().__init__()
         self.feature_map = feature_map
+        if not isinstance(weight, Parameter):
+            self.register_buffer("weight", weight)
         self.weight = weight
         self.bias_module = bias_module
         self.input_transform = input_transform

test/sampling/pathwise/test_prior_samplers.py

@@ -138,16 +138,19 @@ def test_draw_kernel_feature_paths(self):
         with self.subTest("test_initialization"):
             model = self.models[SingleTaskGP][0]
             sample_shape = torch.Size([16])
-            weight_generator = MagicMock()
+            expected_weight_shape = (
+                sample_shape + model.covar_module.batch_shape + (self.num_features,)
+            )
+            weight_generator = MagicMock(
+                side_effect=lambda _: torch.rand(expected_weight_shape)
+            )
             draw_kernel_feature_paths(
                 model=model,
                 sample_shape=sample_shape,
                 num_features=self.num_features,
                 weight_generator=weight_generator,
             )
-            weight_generator.assert_called_once_with(
-                sample_shape + model.covar_module.batch_shape + (self.num_features,)
-            )
+            weight_generator.assert_called_once_with(expected_weight_shape)
 
     def _test_draw_kernel_feature_paths(self, model, paths, sample_shape, atol=3):
         (train_X,) = get_train_inputs(model, transformed=False)