Create dummy habana_frameworks.

HabanaAI · Sep 3, 2024 · c9babcc · c9babcc
1 parent 6ec20c3
commit c9babcc
Showing 1 changed file with 15 additions and 0 deletions.
diff --git a/vllm/cpu_migration.py b/vllm/cpu_migration.py
@@ -1,6 +1,21 @@
+import sys
+import types
+
+# Create dummy habana_frameworks
+habana_frameworks = sys.modules['habana_frameworks'] = types.ModuleType('habana_frameworks')
+torch = sys.modules['habana_frameworks.torch'] = types.ModuleType('habana_frameworks.torch')
+core = sys.modules['habana_frameworks.torch.core'] = types.ModuleType('habana_frameworks.torch.core')
+
+habana_frameworks.torch = torch
+torch.core = core
+core.mark_step = lambda: print('calling mark_step')
+
 import habana_frameworks.torch as htorch
 import torch
 
+# torch.hpu = sys.modules['torch.hpu'] = types.ModuleType('torch.hpu')
+# torch.hpu.synchronize = lambda: print('calling synchronize')
+
 class CpuMigration:
     def __init__(self):
         self._migrate_to_cpu()