# Buggy code
def try_import_torch(error=False):
    """
    Args:
        error (bool): Whether to raise an error if torch cannot be imported.

    Returns:
        tuple: torch AND torch.nn modules.
    """
    if "RLLIB_TEST_NO_TORCH_IMPORT" in os.environ:
        logger.warning("Not importing Torch for test purposes.")
        return None, None

    try:
        import torch
        import torch.nn as nn

        return torch, nn
    except ImportError as e:
        if error:
            raise e
        return None, None


# Traceback (AttributeError)
Traceback (most recent call last):
File "/private/var/tmp/_bazel_travis/7c557718f3877739c657a427203800b1/execroot/com_github_ray_project_ray/bazel-out/darwin-opt/bin/python/ray/tune/test_trial_runner.runfiles/com_github_ray_project_ray/python/ray/tune/tests/test_trial_runner.py", line 6, in <module>
from ray.rllib import _register_all
File "/Users/travis/build/ray-project/ray/python/ray/rllib/__init__.py", line 60, in <module>
_register_all()
File "/Users/travis/build/ray-project/ray/python/ray/rllib/__init__.py", line 37, in _register_all
register_trainable(key, get_agent_class(key))
File "/Users/travis/build/ray-project/ray/python/ray/rllib/agents/registry.py", line 130, in get_agent_class
return _get_agent_class(alg)
File "/Users/travis/build/ray-project/ray/python/ray/rllib/agents/registry.py", line 140, in _get_agent_class
return CONTRIBUTED_ALGORITHMS[alg]()
File "/Users/travis/build/ray-project/ray/python/ray/rllib/contrib/registry.py", line 21, in _import_bandit_lints
from ray.rllib.contrib.bandits.agents.lin_ts import LinTSTrainer
File "/Users/travis/build/ray-project/ray/python/ray/rllib/contrib/bandits/agents/__init__.py", line 1, in <module>
from ray.rllib.contrib.bandits.agents.lin_ts import LinTSTrainer
File "/Users/travis/build/ray-project/ray/python/ray/rllib/contrib/bandits/agents/lin_ts.py", line 5, in <module>
from ray.rllib.contrib.bandits.agents.policy import BanditPolicy
File "/Users/travis/build/ray-project/ray/python/ray/rllib/contrib/bandits/agents/policy.py", line 6, in <module>
from ray.rllib.contrib.bandits.models.linear_regression import \
File "/Users/travis/build/ray-project/ray/python/ray/rllib/contrib/bandits/models/linear_regression.py", line 9, in <module>
class OnlineLinearRegression(nn.Module):
AttributeError: 'NoneType' object has no attribute 'Module'
================================================================================
==================== Test output for //python/ray/tune:test_trial_runner:
/Users/travis/miniconda/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:516: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.
_np_qint8 = np.dtype([("qint8", np.int8, 1)])
/Users/travis/miniconda/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:517: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.
_np_quint8 = np.dtype([("quint8", np.uint8, 1)])
/Users/travis/miniconda/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:518: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.
_np_qint16 = np.dtype([("qint16", np.int16, 1)])
/Users/travis/miniconda/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:519: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.
_np_quint16 = np.dtype([("quint16", np.uint16, 1)])
/Users/travis/miniconda/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:520: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.
_np_qint32 = np.dtype([("qint32", np.int32, 1)])
/Users/travis/miniconda/lib/python3.6/site-packages/tensorflow/python/framework/dtypes.py:525: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.
np_resource = np.dtype([("resource", np.ubyte, 1)])
/Users/travis/miniconda/lib/python3.6/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:541: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.
_np_qint8 = np.dtype([("qint8", np.int8, 1)])
/Users/travis/miniconda/lib/python3.6/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:542: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.
_np_quint8 = np.dtype([("quint8", np.uint8, 1)])
/Users/travis/miniconda/lib/python3.6/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:543: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.
_np_qint16 = np.dtype([("qint16", np.int16, 1)])
/Users/travis/miniconda/lib/python3.6/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:544: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.
_np_quint16 = np.dtype([("quint16", np.uint16, 1)])
/Users/travis/miniconda/lib/python3.6/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:545: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.
_np_qint32 = np.dtype([("qint32", np.int32, 1)])
/Users/travis/miniconda/lib/python3.6/site-packages/tensorboard/compat/tensorflow_stub/dtypes.py:550: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.
np_resource = np.dtype([("resource", np.ubyte, 1)])
lz4 not available, disabling sample compression. This will significantly impact RLlib performance. To install lz4, run `pip install lz4`.
Traceback (most recent call last):
File "/private/var/tmp/_bazel_travis/7c557718f3877739c657a427203800b1/execroot/com_github_ray_project_ray/bazel-out/darwin-opt/bin/python/ray/tune/test_trial_runner.runfiles/com_github_ray_project_ray/python/ray/tune/tests/test_trial_runner.py", line 6, in <module>
from ray.rllib import _register_all
File "/Users/travis/build/ray-project/ray/python/ray/rllib/__init__.py", line 60, in <module>
_register_all()
File "/Users/travis/build/ray-project/ray/python/ray/rllib/__init__.py", line 37, in _register_all
register_trainable(key, get_agent_class(key))
File "/Users/travis/build/ray-project/ray/python/ray/rllib/agents/registry.py", line 130, in get_agent_class
return _get_agent_class(alg)
File "/Users/travis/build/ray-project/ray/python/ray/rllib/agents/registry.py", line 140, in _get_agent_class
return CONTRIBUTED_ALGORITHMS[alg]()
File "/Users/travis/build/ray-project/ray/python/ray/rllib/contrib/registry.py", line 21, in _import_bandit_lints
from ray.rllib.contrib.bandits.agents.lin_ts import LinTSTrainer
File "/Users/travis/build/ray-project/ray/python/ray/rllib/contrib/bandits/agents/__init__.py", line 1, in <module>
from ray.rllib.contrib.bandits.agents.lin_ts import LinTSTrainer
File "/Users/travis/build/ray-project/ray/python/ray/rllib/contrib/bandits/agents/lin_ts.py", line 5, in <module>
from ray.rllib.contrib.bandits.agents.policy import BanditPolicy
File "/Users/travis/build/ray-project/ray/python/ray/rllib/contrib/bandits/agents/policy.py", line 6, in <module>
from ray.rllib.contrib.bandits.models.linear_regression import \
File "/Users/travis/build/ray-project/ray/python/ray/rllib/contrib/bandits/models/linear_regression.py", line 9, in <module>
class OnlineLinearRegression(nn.Module):
AttributeError: 'NoneType' object has no attribute 'Module'

# Fixed code
def try_import_torch(error=False):
    """
    Args:
        error (bool): Whether to raise an error if torch cannot be imported.

    Returns:
        tuple: torch AND torch.nn modules.
    """
    if "RLLIB_TEST_NO_TORCH_IMPORT" in os.environ:
        logger.warning("Not importing Torch for test purposes.")
        return None, None

    try:
        import torch
        import torch.nn as nn

        return torch, nn
    except ImportError as e:
        if error:
            raise e

        nn = NNStub()
        nn.Module = ModuleStub
        return None, nn