[fbsync] Fix typing issue to make DeformConv2d scriptable (#4079)

Reviewed By: fmassa

Differential Revision: D29264317

fbshipit-source-id: 94339e7230eac691ffa5f33f8119ef47caee26b4

test/test_ops.py

@@ -775,6 +775,10 @@ def test_autocast(self, batch_sz, dtype):
         with torch.cuda.amp.autocast():
             self.test_forward(torch.device("cuda"), contiguous=False, batch_sz=batch_sz, dtype=dtype)
 
+    def test_forward_scriptability(self):
+        # Non-regression test for https://github.com/pytorch/vision/issues/4078
+        torch.jit.script(ops.DeformConv2d(in_channels=8, out_channels=8, kernel_size=3))
+
 
 class TestFrozenBNT:
     def test_frozenbatchnorm2d_repr(self):

torchvision/ops/deform_conv.py

@@ -149,7 +149,7 @@ def reset_parameters(self) -> None:
             bound = 1 / math.sqrt(fan_in)
             init.uniform_(self.bias, -bound, bound)
 
-    def forward(self, input: Tensor, offset: Tensor, mask: Tensor = None) -> Tensor:
+    def forward(self, input: Tensor, offset: Tensor, mask: Optional[Tensor] = None) -> Tensor:
         """
         Args:
             input (Tensor[batch_size, in_channels, in_height, in_width]): input tensor