added gpu requirement for this test

huggingface · Dec 18, 2023 · 3c7eb2e · 3c7eb2e
1 parent 96d9ec7
commit 3c7eb2e
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/tests/test_autocast_torchcompatibility_lora.py b/tests/test_autocast_torchcompatibility_lora.py
@@ -7,6 +7,8 @@
 from peft.tuners.lora import Embedding as LoraEmbedding
 from peft.tuners.lora import Linear as LoraLinear
 
+from .testing_utils import require_torch_gpu
+
 
 class SimpleModel(nn.Module):
     def __init__(self):
@@ -117,6 +119,7 @@ def forward(self, input_ids):
         return conv_output
 
 
+@require_torch_gpu
 class TestAutoCast(unittest.TestCase):
     def test_simple_model(self):
         self._test_model(SimpleModel)