-

mila-iqia · Oct 9, 2024 · e780b1d · e780b1d
1 parent 56c3dbd
commit e780b1d
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/benchmarks/lightning/main.py b/benchmarks/lightning/main.py
@@ -1,5 +1,7 @@
 #!/usr/bin/env python
 
+from habana_frameworks.torch import hpu; hpu.init()
+
 import argparse
 import os
 
@@ -55,8 +57,6 @@ def prepare_voir():
     return observer, bench_monitor
 
 def main():
-    from habana_frameworks.torch import hpu; hpu.init()
-
     rank = int(os.getenv("RANK", 0))
     world_size = int(os.getenv("WORLD_SIZE", 1))
     local_world_size = int(os.getenv("LOCAL_WORLD_SIZE", 1))