make indices a buffer in AffineInputTransform (#1656)

Summary:
Pull Request resolved: #1656

see title. This is important for being able to set the device with `to()`

Reviewed By: Balandat, SebastianAment

Differential Revision: D43048429

fbshipit-source-id: 6a290099752b150492162cdc0ff3ed110523ee04

Author: sdaulton

botorch/models/transforms/input.py

@@ -365,7 +365,7 @@ def __init__(
                 raise ValueError("Elements of `indices` have to be smaller than `d`!")
             if len(indices.unique()) != len(indices):
                 raise ValueError("Elements of `indices` tensor must be unique!")
-            self.indices = indices
+            self.register_buffer("indices", indices)
         torch.broadcast_shapes(coefficient.shape, offset.shape)
 
         self._d = d

test/models/transforms/test_input.py

@@ -197,7 +197,13 @@ def test_normalize(self):
             self.assertEqual(nlz.mins.shape, torch.Size([1, 1]))
             self.assertEqual(nlz.ranges.shape, torch.Size([1, 1]))
             self.assertEqual(len(nlz.indices), 1)
-            self.assertTrue((nlz.indices == torch.tensor([0], dtype=torch.long)).all())
+            nlz.to(device=self.device)
+            self.assertTrue(
+                (
+                    nlz.indices
+                    == torch.tensor([0], dtype=torch.long, device=self.device)
+                ).all()
+            )
 
             # test .to
             other_dtype = torch.float if dtype == torch.double else torch.double
@@ -382,17 +388,25 @@ def test_standardize(self):
             self.assertEqual(stdz.means.shape, torch.Size([1, 1]))
             self.assertEqual(stdz.stds.shape, torch.Size([1, 1]))
             self.assertEqual(len(stdz.indices), 1)
+            stdz.to(device=self.device)
             self.assertTrue(
-                torch.equal(stdz.indices, torch.tensor([0], dtype=torch.long))
+                torch.equal(
+                    stdz.indices,
+                    torch.tensor([0], dtype=torch.long, device=self.device),
+                )
             )
             stdz = InputStandardize(d=2, indices=[0], batch_shape=torch.Size([3]))
+            stdz.to(device=self.device)
             self.assertTrue(stdz.training)
             self.assertEqual(stdz._d, 2)
             self.assertEqual(stdz.means.shape, torch.Size([3, 1, 1]))
             self.assertEqual(stdz.stds.shape, torch.Size([3, 1, 1]))
             self.assertEqual(len(stdz.indices), 1)
             self.assertTrue(
-                torch.equal(stdz.indices, torch.tensor([0], dtype=torch.long))
+                torch.equal(
+                    stdz.indices,
+                    torch.tensor([0], device=self.device, dtype=torch.long),
+                )
             )
 
             # test jitter