Fix incompatible type breaking OSS tests (meta-pytorch#1371)

Summary:

Correct typing in dataloader attr to prevent tests from breaking

Differential Revision: D64412835

Author: craymichael

captum/attr/_core/dataloader_attr.py

@@ -3,7 +3,7 @@
 # pyre-strict
 from collections import defaultdict
 from copy import copy
-from typing import Any, Callable, Dict, Iterable, List, Optional, Tuple, Union
+from typing import Any, Callable, cast, Dict, Iterable, List, Optional, Tuple, Union
 
 import torch
 from captum._utils.common import (
@@ -193,8 +193,7 @@ def _forward_with_dataloader(
         feature_mask: Tuple[Tensor, ...],
         # pyre-fixme[24]: Generic type `Callable` expects 2 type parameters.
         reduce: Callable,
-        # pyre-fixme[24]: Generic type `Callable` expects 2 type parameters.
-        to_metric: Optional[Callable],
+        to_metric: Optional[Callable[[Tensor], Tensor]],
         show_progress: bool,
         feature_idx_to_mask_idx: Dict[int, List[int]],
     ) -> Tensor:
@@ -243,7 +242,8 @@ def _forward_with_dataloader(
 
                 accum_states[i] = reduce(accum_states[i], output, perturbed_inputs)
 
-        accum_results = [
+        accum_states = cast(List[Tensor], accum_states)
+        accum_results: List[Tensor] = [
             to_metric(accum) if to_metric else accum for accum in accum_states
         ]
 

captum/attr/_core/lime.py

@@ -522,7 +522,9 @@ def attribute(
             if show_progress:
                 attr_progress.close()
 
-            combined_interp_inps = torch.cat(interpretable_inps).float()
+            # Argument 1 to "cat" has incompatible type "list[Tensor | tuple[Tensor, ...]]";
+            # expected "tuple[Tensor, ...] | list[Tensor]"  [arg-type]
+            combined_interp_inps = torch.cat(interpretable_inps).float()  # type: ignore
             combined_outputs = (
                 torch.cat(outputs)
                 if len(outputs[0].shape) > 0

captum/concept/_utils/classifier.py

@@ -186,7 +186,9 @@ def train_and_eval(
         x_train, x_test, y_train, y_test = _train_test_split(
             torch.cat(inputs), torch.cat(labels), test_split=test_split_ratio
         )
-        self.lm.device = device
+        # error: Incompatible types in assignment (expression has type "str | Any",
+        # variable has type "Tensor | Module")  [assignment]
+        self.lm.device = device  # type: ignore
         self.lm.fit(DataLoader(TensorDataset(x_train, y_train)))
 
         predict = self.lm(x_test)