George-API commited on
Commit
45b968a
·
verified ·
1 Parent(s): f3c357b

Upload run_cloud_training.py with huggingface_hub

Browse files
Files changed (1) hide show
  1. run_cloud_training.py +13 -7
run_cloud_training.py CHANGED
@@ -33,14 +33,20 @@ if 'xformers' in sys.modules:
33
  if 'xformers.ops' in sys.modules:
34
  del sys.modules['xformers.ops']
35
 
36
- # Patch transformers to prevent xformers import
37
- def prevent_xformers_import(name, *args, **kwargs):
38
- if 'xformers' in name:
39
- raise ImportError(f"Import of {name} prevented")
40
- return original_import(name, *args, **kwargs)
 
 
 
 
 
 
41
 
42
- original_import = __import__
43
- __builtins__['__import__'] = prevent_xformers_import
44
 
45
  # Configure PyTorch memory allocator for better memory management
46
  os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
 
33
  if 'xformers.ops' in sys.modules:
34
  del sys.modules['xformers.ops']
35
 
36
+ # Patch Python's import system to prevent xformers from being imported
37
+ class XFormersBlocker:
38
+ def __init__(self, original_importer):
39
+ self.original_importer = original_importer
40
+
41
+ def find_spec(self, fullname, path, target=None):
42
+ if 'xformers' in fullname:
43
+ # Block xformers imports
44
+ return None
45
+ # Use the original importer for everything else
46
+ return self.original_importer.find_spec(fullname, path, target)
47
 
48
+ # Add our import blocker to sys.meta_path
49
+ sys.meta_path.insert(0, XFormersBlocker(sys.meta_path[0]))
50
 
51
  # Configure PyTorch memory allocator for better memory management
52
  os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"