Fix custom modules pyre fix me issues

Differential Revision: D67706756
pytorch · Dec 30, 2024 · 51940e3 · 51940e3
1 parent 7f0d45e
commit 51940e3
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/captum/attr/_utils/custom_modules.py b/captum/attr/_utils/custom_modules.py
@@ -2,6 +2,7 @@
 
 # pyre-strict
 import torch.nn as nn
+from torch import Tensor
 
 
 class Addition_Module(nn.Module):
@@ -12,7 +13,5 @@ class Addition_Module(nn.Module):
     def __init__(self) -> None:
         super().__init__()
 
-    # pyre-fixme[3]: Return type must be annotated.
-    # pyre-fixme[2]: Parameter must be annotated.
-    def forward(self, x1, x2):
+    def forward(self, x1: Tensor, x2: Tensor) -> Tensor:
         return x1 + x2