fix gpu test (keras-team#1939)

* fix gpu test * cast input * update dtype * change to resnet preset * remove arg
ushareng · Oct 24, 2024 · ed035d3 · ed035d3
1 parent f10c45d
commit ed035d3
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 18 deletions.
diff --git a/keras_hub/src/layers/preprocessing/image_converter_test.py b/keras_hub/src/layers/preprocessing/image_converter_test.py
@@ -6,12 +6,10 @@
 from keras import ops
 
 from keras_hub.src.layers.preprocessing.image_converter import ImageConverter
-from keras_hub.src.models.pali_gemma.pali_gemma_backbone import (
-    PaliGemmaBackbone,
-)
 from keras_hub.src.models.pali_gemma.pali_gemma_image_converter import (
     PaliGemmaImageConverter,
 )
+from keras_hub.src.models.resnet.resnet_backbone import ResNetBackbone
 from keras_hub.src.tests.test_case import TestCase
 
 
@@ -86,24 +84,19 @@ def test_from_preset_errors(self):
     def test_save_to_preset(self):
         save_dir = self.get_temp_dir()
         converter = ImageConverter.from_preset(
-            "pali_gemma_3b_mix_224",
+            "resnet_50_imagenet",
             interpolation="nearest",
         )
         converter.save_to_preset(save_dir)
         # Save a tiny backbone so the preset is valid.
-        backbone = PaliGemmaBackbone(
-            vocabulary_size=100,
-            image_size=224,
-            num_layers=1,
-            num_query_heads=1,
-            num_key_value_heads=1,
-            hidden_dim=8,
-            intermediate_dim=16,
-            head_dim=8,
-            vit_patch_size=14,
-            vit_num_heads=1,
-            vit_hidden_dim=8,
-            vit_num_layers=1,
+        backbone = ResNetBackbone(
+            input_conv_filters=[64],
+            input_conv_kernel_sizes=[7],
+            stackwise_num_filters=[64, 64, 64],
+            stackwise_num_blocks=[2, 2, 2],
+            stackwise_num_strides=[1, 2, 2],
+            block_type="basic_block",
+            use_pre_activation=True,
         )
         backbone.save_to_preset(save_dir)
 

diff --git a/keras_hub/src/models/resnet/resnet_backbone.py b/keras_hub/src/models/resnet/resnet_backbone.py
@@ -80,7 +80,6 @@ class ResNetBackbone(FeaturePyramidBackbone):
         stackwise_num_strides=[1, 2, 2],
         block_type="basic_block",
         use_pre_activation=True,
-        pooling="avg",
     )
     model(input_data)
     ```