diff --git a/ultralytics/utils/checks.py b/ultralytics/utils/checks.py index 4143fea4..52f60f89 100644 --- a/ultralytics/utils/checks.py +++ b/ultralytics/utils/checks.py @@ -433,8 +433,8 @@ def check_torchvision(): The compatibility table is a dictionary where the keys are PyTorch versions and the values are lists of compatible Torchvision versions. """ - # Compatibility table compatibility_table = { + "2.6": ["0.21"], "2.5": ["0.20"], "2.4": ["0.19"], "2.3": ["0.18"], @@ -445,7 +445,7 @@ def check_torchvision(): "1.12": ["0.13"], } - # Extract only the major and minor versions + # Check major and minor versions v_torch = ".".join(torch.__version__.split("+")[0].split(".")[:2]) if v_torch in compatibility_table: compatible_versions = compatibility_table[v_torch]