Merge pull request #45 from mit-acl/feature/torch-support

add patch to support torch>2.4

Author: liqyn

roman/map/fastsam_wrapper.py

@@ -38,6 +38,17 @@
 logger = logging.getLogger(__name__)
 logger.setLevel(logging.WARN)
 
+# Patch torch.load to disable weights_only loading for torch>2.4
+torch_version = torch.__version__.split('.')
+if int(torch_version[0]) > 2 or (int(torch_version[0]) == 2 and int(torch_version[1]) > 4):
+    _real_torch_load = torch.load
+
+    def torch_load_no_weights_only(*args, **kwargs):
+        kwargs["weights_only"] = False
+        return _real_torch_load(*args, **kwargs)
+
+    torch.load = torch_load_no_weights_only
+
 class FastSAMWrapper():
 
     def __init__(self, 

setup.py

@@ -27,8 +27,8 @@
         'shapely',
         'opencv-python>=4.6.0',
         'pyyaml',
-        'torch==2.4.0',
-        'torchvision==0.19.0',
+        'torch>=2.4.0',
+        'torchvision>=0.19.0',
         'fastsam @ git+https://github.com/CASIA-IVA-Lab/FastSAM.git@4d153e9',
         'robotdatapy>=1.1.0',
         'gdown',