Skip to content

Commit

Permalink
[run-slow] Fixing minimum
Browse files Browse the repository at this point in the history
  • Loading branch information
qubvel committed Dec 23, 2024
1 parent 15b0658 commit 238dd11
Show file tree
Hide file tree
Showing 3 changed files with 5 additions and 3 deletions.
3 changes: 2 additions & 1 deletion tests/models/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -137,6 +137,7 @@ def test_classification_head(self):

self.assertEqual(cls_probs.shape[1], 10)

@requires_torch_greater_or_equal("2.0.0")
def test_save_load_with_hub_mixin(self):
# instantiate model
model = smp.create_model(
Expand Down Expand Up @@ -172,7 +173,7 @@ def test_save_load_with_hub_mixin(self):
self.assertIn("my_awesome_metric", readme)

@slow_test
@requires_torch_greater_or_equal("2.0.1")
@requires_torch_greater_or_equal("2.0.0")
def test_preserve_forward_output(self):
from huggingface_hub import hf_hub_download

Expand Down
3 changes: 2 additions & 1 deletion tests/models/test_segformer.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,14 +3,15 @@
import segmentation_models_pytorch as smp

from tests.models import base
from tests.utils import slow_test, default_device
from tests.utils import slow_test, default_device, requires_torch_greater_or_equal


@pytest.mark.segformer
class TestSegformerModel(base.BaseModelTester):
test_model_type = "segformer"

@slow_test
@requires_torch_greater_or_equal("2.0.0")
def test_load_pretrained(self):
hub_checkpoint = "smp-hub/segformer-b0-512x512-ade-160k"

Expand Down
2 changes: 1 addition & 1 deletion tests/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,6 @@ def requires_torch_greater_or_equal(version: str):
torch_version = Version(torch.__version__)
provided_version = Version(version)
return unittest.skipUnless(
torch_version >= provided_version,
torch_version < provided_version,
f"torch version {torch_version} is less than {provided_version}",
)

0 comments on commit 238dd11

Please sign in to comment.