diff --git a/requirements.txt b/requirements.txt index 0e6e860a5a..68cd4298f0 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,5 +1,3 @@ -torch -numpy sentencepiece packaging expecttest # So we can use IS_FBCODE flag diff --git a/setup.py b/setup.py index d4cf988b43..fc14725005 100644 --- a/setup.py +++ b/setup.py @@ -88,6 +88,16 @@ def get_extensions(): return ext_modules +# Mimic code from torchvision https://github.com/pytorch/vision/blob/143d078b28f00471156a4e562dd3836370acc9ee/setup.py#L58 +pytorch_dep = "torch" +if os.getenv("PYTORCH_VERSION"): + pytorch_dep += "==" + os.getenv("PYTORCH_VERSION") + +requirements = [ + "numpy", + pytorch_dep, +] + setup( name=package_name, version=version+version_suffix, @@ -97,7 +107,7 @@ def get_extensions(): "torchao.kernel.configs": ["*.pkl"], }, ext_modules=get_extensions() if use_cpp != "0" else None, - install_requires=read_requirements("requirements.txt"), + install_requires=requirements, extras_require={"dev": read_requirements("dev-requirements.txt")}, description="Package for applying ao techniques to GPU models", long_description=open("README.md").read(),