From 13655728cda68ce4d8eefb92124b3b2191991dce Mon Sep 17 00:00:00 2001 From: Yanhui Liang Date: Wed, 12 Feb 2020 16:15:52 -0800 Subject: [PATCH] Fix the timeout of nasnet test. PiperOrigin-RevId: 294782893 Change-Id: Iffc97ba7ecb2072fd6d42ba7d9923952b157452d --- tensorflow/python/keras/applications/BUILD | 21 ++++++++++++++++--- .../applications_load_weight_test.py | 3 ++- 2 files changed, 20 insertions(+), 4 deletions(-) diff --git a/tensorflow/python/keras/applications/BUILD b/tensorflow/python/keras/applications/BUILD index 9457e4e1272..1dc0397b2fb 100644 --- a/tensorflow/python/keras/applications/BUILD +++ b/tensorflow/python/keras/applications/BUILD @@ -235,14 +235,29 @@ tf_py_test( ) tf_py_test( - name = "applications_load_weight_test_nasnet", + name = "applications_load_weight_test_nasnet_mobile", srcs = ["applications_load_weight_test.py"], - args = ["--module=nasnet"], + args = ["--module=nasnet_mobile"], + main = "applications_load_weight_test.py", + tags = [ + "no_oss", + "no_pip", + ], + deps = [ + ":applications", + "//tensorflow/python:client_testlib", + "@absl_py//absl/testing:parameterized", + ], +) + +tf_py_test( + name = "applications_load_weight_test_nasnet_large", + srcs = ["applications_load_weight_test.py"], + args = ["--module=nasnet_large"], main = "applications_load_weight_test.py", tags = [ "no_oss", "no_pip", - "notsan", # b/149100891 ], deps = [ ":applications", diff --git a/tensorflow/python/keras/applications/applications_load_weight_test.py b/tensorflow/python/keras/applications/applications_load_weight_test.py index d33e844981b..42146c66f97 100644 --- a/tensorflow/python/keras/applications/applications_load_weight_test.py +++ b/tensorflow/python/keras/applications/applications_load_weight_test.py @@ -53,7 +53,8 @@ ARG_TO_MODEL = { 'mobilenet_v2': (mobilenet_v2, [mobilenet_v2.MobileNetV2]), 'densenet': (densenet, [densenet.DenseNet121, densenet.DenseNet169, densenet.DenseNet201]), - 'nasnet': (nasnet, [nasnet.NASNetMobile, nasnet.NASNetLarge]), + 'nasnet_mobile': (nasnet, [nasnet.NASNetMobile]), + 'nasnet_large': (nasnet, [nasnet.NASNetLarge]), 'efficientnet': (efficientnet, [efficientnet.EfficientNetB0, efficientnet.EfficientNetB1, efficientnet.EfficientNetB2, efficientnet.EfficientNetB3,