ray/rllib/tests/test_dependency_torch.py
Balaji Veeramani 7f1bacc7dc
[CI] Format Python code with Black (#21975)
See #21316 and #21311 for the motivation behind these changes.
2022-01-29 18:41:57 -08:00

36 lines
978 B
Python
Executable file

#!/usr/bin/env python
import os
import sys
if __name__ == "__main__":
# Do not import torch for testing purposes.
os.environ["RLLIB_TEST_NO_TORCH_IMPORT"] = "1"
from ray.rllib.agents.a3c import A2CTrainer
assert "torch" not in sys.modules, "`torch` initially present, when it shouldn't!"
# Note: No ray.init(), to test it works without Ray
trainer = A2CTrainer(
env="CartPole-v0",
config={
"framework": "tf",
"num_workers": 0,
# Disable the logger due to a sort-import attempt of torch
# inside the tensorboardX.SummaryWriter class.
"logger_config": {
"type": "ray.tune.logger.NoopLogger",
},
},
)
trainer.train()
assert "torch" not in sys.modules, (
"`torch` should not be imported after creating and " "training A3CTrainer!"
)
# Clean up.
del os.environ["RLLIB_TEST_NO_TORCH_IMPORT"]
print("ok")