Fix neuron conductance pyre fixme issues

Differential Revision: D67523217

Author: Vivek Miglani

captum/attr/_core/neuron/neuron_conductance.py

@@ -39,8 +39,7 @@ class NeuronConductance(NeuronAttribution, GradientAttribution):
 
     def __init__(
         self,
-        # pyre-fixme[24]: Generic type `Callable` expects 2 type parameters.
-        forward_func: Callable,
+        forward_func: Callable[..., Tensor],
         layer: Module,
         device_ids: Union[None, List[int]] = None,
         multiply_by_inputs: bool = True,
@@ -94,8 +93,9 @@ def __init__(
     def attribute(
         self,
         inputs: TensorOrTupleOfTensorsGeneric,
-        # pyre-fixme[24]: Generic type `Callable` expects 2 type parameters.
-        neuron_selector: Union[int, Tuple[int, ...], Callable],
+        neuron_selector: Union[
+            int, Tuple[int, ...], Callable[[Union[Tensor, Tuple[Tensor, ...]]], Tensor]
+        ],
         baselines: BaselineType = None,
         target: TargetType = None,
         additional_forward_args: Optional[object] = None,
@@ -285,8 +285,6 @@ def attribute(
                 " results.",
                 stacklevel=1,
             )
-        # pyre-fixme[6]: For 1st argument expected `Tensor` but got
-        #  `TensorOrTupleOfTensorsGeneric`.
         is_inputs_tuple = _is_tuple(inputs)
 
         # pyre-fixme[9]: inputs has type `TensorOrTupleOfTensorsGeneric`; used as
@@ -334,8 +332,9 @@ def attribute(
     def _attribute(
         self,
         inputs: Tuple[Tensor, ...],
-        # pyre-fixme[24]: Generic type `Callable` expects 2 type parameters.
-        neuron_selector: Union[int, Tuple[int, ...], Callable],
+        neuron_selector: Union[
+            int, Tuple[int, ...], Callable[[Union[Tensor, Tuple[Tensor, ...]]], Tensor]
+        ],
         baselines: Tuple[Union[Tensor, int, float], ...],
         target: TargetType = None,
         additional_forward_args: Optional[object] = None,