Check shape of state dict when comparing input transforms (#3051)

Summary:
## Motivation

See #3050

### Have you read the [Contributing Guidelines on pull requests](https://github.com/meta-pytorch/botorch/blob/main/CONTRIBUTING.md#pull-requests)?

Y

Pull Request resolved: #3051

Test Plan: I have added an additional test case to `test/models/transforms/test_input.py`. Without the changes in this PR, the new test case would fail with a runtime error, as described in #3050. The test now correctly passes.

Reviewed By: saitcakmak

Differential Revision: D85149051

Pulled By: Balandat

fbshipit-source-id: 4f6a972601994154d19edd44fef43ddfa092170f

Author: TobyBoyne

botorch/models/transforms/input.py

@@ -42,6 +42,13 @@
 from torch.nn.functional import one_hot
 
 
+def _allclose(input: Tensor, other: Tensor) -> bool:
+    """Check if `input` and `other` are the same shape, and satisfy `torch.allclose`."""
+    if input.shape != other.shape:
+        return False
+    return torch.allclose(input, other)
+
+
 class InputTransform(Module, ABC):
     r"""Abstract base class for input transforms.
 
@@ -124,7 +131,7 @@ def equals(self, other: InputTransform) -> bool:
             and (self.transform_on_eval == other.transform_on_eval)
             and (self.transform_on_fantasize == other.transform_on_fantasize)
             and all(
-                torch.allclose(v, other_state_dict[k].to(v))
+                _allclose(v, other_state_dict[k].to(v))
                 for k, v in self.state_dict().items()
             )
         )

test/models/transforms/test_input.py

@@ -400,7 +400,9 @@ def test_normalize(self) -> None:
                 self.assertTrue(nlz6.equals(nlz8))
                 nlz9 = Normalize(d=3, batch_shape=batch_shape, indices=[0, 1])
                 nlz10 = Normalize(d=3, batch_shape=batch_shape, indices=[0, 2])
+                nlz11 = Normalize(d=3, batch_shape=batch_shape, indices=[0, 1, 2])
                 self.assertFalse(nlz9.equals(nlz10))
+                self.assertFalse(nlz9.equals(nlz11))
 
                 # test with grad
                 nlz = Normalize(d=1)