Update experience_source.py #749 Bug Fix (#751)

* Update experience_source.py * Update changelog Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Aki Nitta <nitta@akihironitta.com>
Lightning-Universe · Oct 15, 2021 · f4f6d53 · f4f6d53
1 parent d4d1757
commit f4f6d53
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -21,7 +21,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Fixed
 
-
+- Fix doctest fails with ImportError: cannot import name 'Env' from 'gym' ([#751](https://github.com/PyTorchLightning/lightning-bolts/pull/751))
 
 ## [0.4.0] - 2021-09-09
 

diff --git a/pl_bolts/datamodules/experience_source.py b/pl_bolts/datamodules/experience_source.py
@@ -10,9 +10,13 @@
 from pl_bolts.utils import _GYM_AVAILABLE
 from pl_bolts.utils.warnings import warn_missing_pkg
 
-if _GYM_AVAILABLE:
-    from gym import Env
-else:  # pragma: no cover
+try:
+    if _GYM_AVAILABLE:
+        from gym import Env
+    else:  # pragma: no cover
+        warn_missing_pkg("gym")
+        Env = object
+except:
     warn_missing_pkg("gym")
     Env = object