use no_grad when using input transform in model.__init__ (#610)

Summary:
Pull Request resolved: #610

See title.

Reviewed By: Balandat

Differential Revision: D25158547

fbshipit-source-id: 393c5b6ba45909c13cf876893d28ca9bbcd7e83e

Author: sdaulton

botorch/models/gp_regression.py

@@ -92,9 +92,10 @@ def __init__(
         """
         if input_transform is not None:
             input_transform.to(train_X)
-        transformed_X = self.transform_inputs(
-            X=train_X, input_transform=input_transform
-        )
+        with torch.no_grad():
+            transformed_X = self.transform_inputs(
+                X=train_X, input_transform=input_transform
+            )
         if outcome_transform is not None:
             train_Y, _ = outcome_transform(train_Y)
         self._validate_tensor_args(X=transformed_X, Y=train_Y)
@@ -206,9 +207,10 @@ def __init__(
         """
         if input_transform is not None:
             input_transform.to(train_X)
-        transformed_X = self.transform_inputs(
-            X=train_X, input_transform=input_transform
-        )
+        with torch.no_grad():
+            transformed_X = self.transform_inputs(
+                X=train_X, input_transform=input_transform
+            )
         if outcome_transform is not None:
             train_Y, train_Yvar = outcome_transform(train_Y, train_Yvar)
         self._validate_tensor_args(X=transformed_X, Y=train_Y, Yvar=train_Yvar)

botorch/models/gp_regression_fidelity.py

@@ -97,9 +97,10 @@ def __init__(
             )
         if input_transform is not None:
             input_transform.to(train_X)
-        transformed_X = self.transform_inputs(
-            X=train_X, input_transform=input_transform
-        )
+        with torch.no_grad():
+            transformed_X = self.transform_inputs(
+                X=train_X, input_transform=input_transform
+            )
 
         self._set_dimensions(train_X=transformed_X, train_Y=train_Y)
         covar_module, subset_batch_dict = _setup_multifidelity_covar_module(
@@ -208,9 +209,10 @@ def __init__(
             )
         if input_transform is not None:
             input_transform.to(train_X)
-        transformed_X = self.transform_inputs(
-            X=train_X, input_transform=input_transform
-        )
+        with torch.no_grad():
+            transformed_X = self.transform_inputs(
+                X=train_X, input_transform=input_transform
+            )
         self._set_dimensions(train_X=transformed_X, train_Y=train_Y)
         covar_module, subset_batch_dict = _setup_multifidelity_covar_module(
             dim=transformed_X.size(-1),

botorch/models/multitask.py

@@ -83,9 +83,10 @@ def __init__(
         """
         if input_transform is not None:
             input_transform.to(train_X)
-        transformed_X = self.transform_inputs(
-            X=train_X, input_transform=input_transform
-        )
+        with torch.no_grad():
+            transformed_X = self.transform_inputs(
+                X=train_X, input_transform=input_transform
+            )
         self._validate_tensor_args(X=transformed_X, Y=train_Y)
         all_tasks, task_feature, d = self.get_all_tasks(
             transformed_X, task_feature, output_tasks
@@ -292,9 +293,10 @@ def __init__(
         """
         if input_transform is not None:
             input_transform.to(train_X)
-        transformed_X = self.transform_inputs(
-            X=train_X, input_transform=input_transform
-        )
+        with torch.no_grad():
+            transformed_X = self.transform_inputs(
+                X=train_X, input_transform=input_transform
+            )
         self._validate_tensor_args(X=transformed_X, Y=train_Y, Yvar=train_Yvar)
         # We'll instatiate a MultiTaskGP and simply override the likelihood
         super().__init__(