Prepare for "Fix type-safety of torch.nn.Module instances": fbcode/p*

captum/attr/_core/feature_ablation.py

@@ -808,7 +808,10 @@ def _construct_ablated_input(
         current_mask = current_mask.to(expanded_input.device)
         assert baseline is not None, "baseline must be provided"
         ablated_tensor = (
-            expanded_input * (1 - current_mask).to(expanded_input.dtype)
+            expanded_input
+            * (1 - current_mask).to(expanded_input.dtype)
+            # pyre-fixme[58]: `*` is not supported for operand types `Union[None, float,
+            #  Tensor]` and `Tensor`.
         ) + (baseline * current_mask.to(expanded_input.dtype))
         return ablated_tensor, current_mask
 

captum/attr/_core/occlusion.py

@@ -323,6 +323,8 @@ def _construct_ablated_input(
                 torch.ones(1, dtype=torch.long, device=expanded_input.device)
                 - input_mask
             ).to(expanded_input.dtype)
+            # pyre-fixme[58]: `*` is not supported for operand types `Union[None, float,
+            #  Tensor]` and `Tensor`.
         ) + (baseline * input_mask.to(expanded_input.dtype))
         return ablated_tensor, input_mask