use pytest.mark directly (#27390)

fix

Co-authored-by: ydshieh <ydshieh@users.noreply.github.com>
This commit is contained in:
Yih-Dar
2023-11-09 13:32:54 +01:00
committed by GitHub
parent 791ec370d1
commit 3258ff9330
5 changed files with 14 additions and 14 deletions

View File

@@ -16,7 +16,7 @@ import os
import tempfile
import unittest
from pytest import mark
import pytest
from transformers import DistilBertConfig, is_torch_available
from transformers.testing_utils import require_flash_attn, require_torch, require_torch_accelerator, slow, torch_device
@@ -290,7 +290,7 @@ class DistilBertModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCa
# Because DistilBertForMultipleChoice requires inputs with different shapes we need to override this test.
@require_flash_attn
@require_torch_accelerator
@mark.flash_attn_test
@pytest.mark.flash_attn_test
@slow
def test_flash_attn_2_inference(self):
import torch
@@ -344,7 +344,7 @@ class DistilBertModelTest(ModelTesterMixin, PipelineTesterMixin, unittest.TestCa
# Because DistilBertForMultipleChoice requires inputs with different shapes we need to override this test.
@require_flash_attn
@require_torch_accelerator
@mark.flash_attn_test
@pytest.mark.flash_attn_test
@slow
def test_flash_attn_2_inference_padding_right(self):
import torch