Skip to content

Commit

Permalink
[testing] add timeout_decorator (huggingface#3543)
Browse files Browse the repository at this point in the history
  • Loading branch information
sshleifer authored May 1, 2020
1 parent b868617 commit 18db92d
Show file tree
Hide file tree
Showing 3 changed files with 5 additions and 2 deletions.
1 change: 1 addition & 0 deletions setup.cfg
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ known_third_party =
tensorboardX
tensorflow
tensorflow_datasets
timeout_decorator
torch
torchtext
torchvision
Expand Down
2 changes: 1 addition & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -72,7 +72,7 @@
extras["serving"] = ["pydantic", "uvicorn", "fastapi", "starlette"]
extras["all"] = extras["serving"] + ["tensorflow", "torch"]

extras["testing"] = ["pytest", "pytest-xdist"]
extras["testing"] = ["pytest", "pytest-xdist", "timeout-decorator"]
extras["docs"] = ["recommonmark", "sphinx", "sphinx-markdown-tables", "sphinx-rtd-theme"]
extras["quality"] = [
"black",
Expand Down
4 changes: 3 additions & 1 deletion tests/test_modeling_bart.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,10 +13,11 @@
# See the License for the specific language governing permissions and
# limitations under the License.


import tempfile
import unittest

import timeout_decorator # noqa

from transformers import is_torch_available

from .test_configuration_common import ConfigTester
Expand Down Expand Up @@ -357,6 +358,7 @@ def test_sequence_classification_forward(self):
loss = outputs[0]
self.assertIsInstance(loss.item(), float)

@timeout_decorator.timeout(1)
def test_lm_forward(self):
config, input_ids, batch_size = self._get_config_and_data()
lm_labels = ids_tensor([batch_size, input_ids.shape[1]], self.vocab_size).to(torch_device)
Expand Down

0 comments on commit 18db92d

Please sign in to comment.