Enable torch tracing by changing assertions in d2go forwards to allow for...

Enable torch tracing by changing assertions in d2go forwards to allow for torch.fx.proxy.Proxy type.

Summary:
X-link: https://github.com/facebookresearch/detectron2/pull/4227

Pull Request resolved: https://github.com/facebookresearch/d2go/pull/241

Torch FX tracing propagates a type of `torch.fx.proxy.Proxy` through the graph.

Existing type assertions in the d2go code base trigger during torch FX tracing, causing tracing to fail.

This adds a check for FX tracing in progress and  adds a helper function `assert_fx_safe()`, that can be used in place of a standard assertion. This function only applies the assertion if one is not tracing, allowing d2go assertion tests to be compatible with FX tracing.

Reviewed By: wat3rBro

Differential Revision: D35518556

fbshipit-source-id: a9b5d3d580518ca74948544973ae89f8b9de3282

d2go/modeling/backbone/modules.py

@@ -7,6 +7,7 @@ from typing import List
 import torch
 import torch.nn as nn
 from detectron2 import layers
+from detectron2.utils.tracing import assert_fx_safe
 from mobile_cv.arch.fbnet_v2.irf_block import IRFBlock
 
 
@@ -33,7 +34,7 @@ class RPNHeadConvRegressor(nn.Module):
             torch.nn.init.constant_(l.bias, 0)
 
     def forward(self, x: List[torch.Tensor]):
-        assert isinstance(x, (list, tuple))
+        assert_fx_safe(isinstance(x, (list, tuple)), "Unexpected data type")
         logits = [self.cls_logits(y) for y in x]
         bbox_reg = [self.bbox_pred(y) for y in x]