update tests/test_models

Author: brianhou0208

tests/test_models.py

@@ -28,6 +28,7 @@ def get_sample(model_class):
         smp.Unet,
         smp.UnetPlusPlus,
         smp.MAnet,
+        smp.Segformer,
     ]:
         sample = torch.ones([1, 3, 64, 64])
     elif model_class == smp.PAN:
@@ -61,7 +62,16 @@ def _test_forward_backward(model, sample, test_shape=False):
 @pytest.mark.parametrize("encoder_depth", [3, 5])
 @pytest.mark.parametrize(
     "model_class",
-    [smp.FPN, smp.PSPNet, smp.Linknet, smp.Unet, smp.UnetPlusPlus, smp.UPerNet],
+    [
+        smp.FPN,
+        smp.PSPNet,
+        smp.Linknet,
+        smp.Unet,
+        smp.UnetPlusPlus,
+        smp.MAnet,
+        smp.UPerNet,
+        smp.Segformer,
+    ],
 )
 def test_forward(model_class, encoder_name, encoder_depth, **kwargs):
     if (
@@ -106,6 +116,7 @@ def test_forward(model_class, encoder_name, encoder_depth, **kwargs):
         smp.DeepLabV3,
         smp.DeepLabV3Plus,
         smp.UPerNet,
+        smp.Segformer,
     ],
 )
 def test_forward_backward(model_class):
@@ -127,6 +138,7 @@ def test_forward_backward(model_class):
         smp.DeepLabV3,
         smp.DeepLabV3Plus,
         smp.UPerNet,
+        smp.Segformer,
     ],
 )
 def test_aux_output(model_class):