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

…p* (pytorch#1448) Summary: See D52890934 Differential Revision: D66235323
ezyang · Nov 20, 2024 · 349d3f7 · 349d3f7
1 parent 72a32af
commit 349d3f7
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/captum/attr/_core/feature_ablation.py b/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
 

diff --git a/captum/attr/_core/occlusion.py b/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