Fix issues with CI builds

* Relax version of bitsandbytes * Add triton to requirements * Few fixes for using newer version of torch
microsoft · Jan 31, 2025 · e1e2810 · e1e2810
1 parent d98186d
commit e1e2810
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 4 deletions.
diff --git a/olive/model/handler/pytorch.py b/olive/model/handler/pytorch.py
@@ -151,7 +151,7 @@ def load_model(self, rank: int = None, cache_model: bool = True) -> "torch.nn.Mo
             elif self.model_file_format == ModelFileFormat.PYTORCH_TORCH_SCRIPT:
                 model = torch.jit.load(self.model_path)
             elif self.model_file_format == ModelFileFormat.PYTORCH_ENTIRE_MODEL:
-                model = torch.load(self.model_path)
+                model = torch.load(self.model_path, weights_only=False)
             elif self.model_file_format == ModelFileFormat.PYTORCH_SLICE_GPT_MODEL:
                 model = self._load_slicegpt_model()
             elif self.model_file_format == ModelFileFormat.PYTORCH_STATE_DICT:

diff --git a/test/requirements-test-gpu.txt b/test/requirements-test-gpu.txt
@@ -1,5 +1,5 @@
 -r requirements-test.txt
 auto-gptq
 autoawq
-# only available on Linux currently
-bitsandbytes==0.43.3
+bitsandbytes
+triton
diff --git a/test/unit_test/model/test_pytorch_model.py b/test/unit_test/model/test_pytorch_model.py
@@ -58,7 +58,7 @@ def test_load_from_path(torch_load):
     model = PyTorchModelHandler(model_path="test_path")
 
     assert model.load_model() == "dummy_pytorch_model"
-    torch_load.assert_called_once_with("test_path")
+    torch_load.assert_called_once_with("test_path", weights_only=False)
 
 
 @patch("olive.model.handler.pytorch.UserModuleLoader")