Add check for fx compatibility on segmentation and video models (#4131)

* Add check for fx compatibility on segmentation models.

* Add fx check on video models.

test/test_models.py

@@ -466,6 +466,7 @@ def test_segmentation_model(model_name, dev):
     full_validation = check_out(out)
 
     _check_jit_scriptable(model, (x,), unwrapper=script_model_unwrapper.get(model_name, None))
+    _check_fx_compatible(model, x)
 
     if dev == torch.device("cuda"):
         with torch.cuda.amp.autocast():
@@ -616,6 +617,7 @@ def test_video_model(model_name, dev):
     x = torch.rand(input_shape).to(device=dev)
     out = model(x)
     _check_jit_scriptable(model, (x,), unwrapper=script_model_unwrapper.get(model_name, None))
+    _check_fx_compatible(model, x)
     assert out.shape[-1] == 50
 
     if dev == torch.device("cuda"):