added num_classses arg

keras-team · Sep 6, 2023 · 118d392 · 118d392
1 parent 2bcb664
commit 118d392
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/keras_cv/models/segmentation/deeplab_v3_plus/deeplab_v3_plus_test.py b/keras_cv/models/segmentation/deeplab_v3_plus/deeplab_v3_plus_test.py
@@ -74,9 +74,11 @@ def test_weights_change(self):
             self.assertNotAllEqual(w1, w2)
             self.assertFalse(ops.any(ops.isnan(w2)))
 
-    @pytest.mark.large
+    # @pytest.mark.large
     def test_with_model_preset(self):
-        model = DeepLabV3Plus.from_preset("deeplab_v3_plus_resnet50_pascalvoc")
+        model = DeepLabV3Plus.from_preset(
+            "deeplab_v3_plus_resnet50_pascalvoc", num_classes=21,
+        )
         images = np.random.uniform(size=(2, 512, 512, 3))
         _ = model(images)
         _ = model.predict(images)