From b10b11deb0ab6146f9173465be758238deb13219 Mon Sep 17 00:00:00 2001
From: Jirka Borovec <Borda@users.noreply.github.com>
Date: Fri, 20 Nov 2020 21:32:13 +0100
Subject: [PATCH] add version for CPU users (#4794)

---
 README.md |  1 +
 setup.py  | 10 ++++++++++
 2 files changed, 11 insertions(+)

diff --git a/README.md b/README.md
index e7ecfc1a45..91e9b6e73a 100644
--- a/README.md
+++ b/README.md
@@ -115,6 +115,7 @@ Simple installation from PyPI
 ```bash
 pip install pytorch-lightning
 ```
+_To get full package experience you can install also all optional dependencies with `pytorch-lightning['extra']` or for CPU users with `pytorch-lightning['cpu-extra']`._
 
 From Conda
 ```bash
diff --git a/setup.py b/setup.py
index e811f7a33d..fbbf61142b 100755
--- a/setup.py
+++ b/setup.py
@@ -73,6 +73,16 @@ extras = {
 extras['dev'] = extras['extra'] + extras['loggers'] + extras['test']
 extras['all'] = extras['dev'] + extras['examples']  # + extras['docs']
 
+# These packages shall be installed only on GPU machines
+PACKAGES_GPU_ONLY = (
+    'horovod',
+)
+# create a version for CPU machines
+for ex in ('cpu', 'cpu-extra'):
+    kw = ex.split('-')[1] if '-' in ex else 'all'
+    # filter cpu only packages
+    extras[ex] = [pkg for pkg in extras[kw] if not any(pgpu.lower() in pkg.lower() for pgpu in PACKAGES_GPU_ONLY)]
+
 # https://packaging.python.org/discussions/install-requires-vs-requirements /
 # keep the meta-data here for simplicity in reading this file... it's not obvious
 # what happens and to non-engineers they won't know to look in init ...