FIX Avoid early import of torch extension by BOFT (#1879)

huggingface · Jun 26, 2024 · 7ffa43b · 7ffa43b
1 parent 27bc305
commit 7ffa43b
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/peft/tuners/boft/layer.py b/src/peft/tuners/boft/layer.py
@@ -27,7 +27,6 @@
 import torch.nn as nn
 import torch.nn.functional as F
 from torch.autograd import Function
-from torch.utils.cpp_extension import load
 
 from peft.tuners.tuners_utils import BaseTunerLayer, check_adapters_to_merge
 
@@ -78,6 +77,9 @@ def get_fbd_cuda():
     if _FBD_CUDA is not None:
         return _FBD_CUDA
 
+    # This import initializes cuda context and should thus be local, see issue 1877
+    from torch.utils.cpp_extension import load
+
     curr_dir = os.path.dirname(__file__)
     # need ninja to build the extension
     try: