Unverified Commit 5976fe19 authored by Tejaswin Parthasarathy's avatar Tejaswin Parthasarathy Committed by GitHub
Browse files

Fix TE discovery in python virtual environments (#303)



fix : TE virtuelenv discovery
Signed-off-by: default avatartejaswinp <tejaswinp@nvidia.com>
parent 6f7853dc
...@@ -7,15 +7,16 @@ import ctypes ...@@ -7,15 +7,16 @@ import ctypes
import os import os
import platform import platform
import subprocess import subprocess
import sys
def get_te_path(): def get_te_path():
"""Find Transformer Engine install path using pip""" """Find Transformer Engine install path using pip"""
command = ["pip", "show", "transformer_engine"] command = [sys.executable, "-m", "pip", "show", "transformer_engine"]
result = subprocess.run(command, capture_output=True, check=True, text=True) result = subprocess.run(command, capture_output=True, check=True, text=True)
result = result.stdout.replace("\n", ":").split(":") result = result.stdout.replace("\n", ":").split(":")
return result[result.index("Location")+1].strip() return result[result.index("Location") + 1].strip()
def _load_library(): def _load_library():
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment