Fix sensitivity pyre fixme issues (#1599)

Summary:
Pull Request resolved: #1599

Fixing unresolved pyre fixme issues in corresponding file

Differential Revision: D76737670

Author: Vivek Miglani

captum/metrics/_core/sensitivity.py

@@ -47,15 +47,13 @@ def default_perturb_func(
                 original inputs.
 
     """
-    # pyre-fixme[9]: inputs has type `TensorOrTupleOfTensorsGeneric`; used as
-    #  `Tuple[Tensor, ...]`.
-    inputs = _format_tensor_into_tuples(inputs)
+    inputs_tuple = _format_tensor_into_tuples(inputs)
     perturbed_input = tuple(
         input
         + torch.FloatTensor(input.size())  # type: ignore
         .uniform_(-perturb_radius, perturb_radius)
         .to(input.device)
-        for input in inputs
+        for input in inputs_tuple
     )
     return perturbed_input
 
@@ -64,8 +62,9 @@ def default_perturb_func(
 def sensitivity_max(
     explanation_func: Callable[..., TensorOrTupleOfTensorsGeneric],
     inputs: TensorOrTupleOfTensorsGeneric,
-    # pyre-fixme[24]: Generic type `Callable` expects 2 type parameters.
-    perturb_func: Callable = default_perturb_func,
+    perturb_func: Callable[
+        ..., Union[Tensor, Tuple[Tensor, ...]]
+    ] = default_perturb_func,
     perturb_radius: float = 0.02,
     n_perturb_samples: int = 10,
     norm_ord: str = "fro",
@@ -202,7 +201,7 @@ def sensitivity_max(
 
     def _generate_perturbations(
         current_n_perturb_samples: int,
-    ) -> TensorOrTupleOfTensorsGeneric:
+    ) -> Union[Tensor, Tuple[Tensor, ...]]:
         r"""
         The perturbations are generated for each example
         `current_n_perturb_samples` times.
@@ -228,8 +227,7 @@ def max_values(input_tnsr: Tensor) -> Tensor:
         return torch.max(input_tnsr, dim=1).values  # type: ignore
 
     kwarg_expanded_for = None
-    # pyre-fixme[33]: Given annotation cannot be `Any`.
-    kwargs_copy: Any = None
+    kwargs_copy: object = None
 
     def _next_sensitivity_max(current_n_perturb_samples: int) -> Tensor:
         inputs_perturbed = _generate_perturbations(current_n_perturb_samples)
@@ -254,9 +252,7 @@ def _next_sensitivity_max(current_n_perturb_samples: int) -> Tensor:
                 )
                 if (
                     isinstance(baselines[0], Tensor)
-                    # pyre-fixme[16]: Item `float` of `Union[float, int, Tensor]`
-                    #  has no attribute `shape`.
-                    and baselines[0].shape == inputs[0].shape
+                    and cast(Tensor, baselines[0]).shape == inputs[0].shape
                 ):
                     _expand_and_update_baselines(
                         cast(Tuple[Tensor, ...], inputs),