From 75592e664f3283631f9ff3365c033974fc493013 Mon Sep 17 00:00:00 2001 From: Sven Mika Date: Tue, 21 Jul 2020 22:02:24 +0200 Subject: [PATCH] Issue 9568: `rllib train` framework in config gets overridden with tf. (#9572) --- rllib/train.py | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/rllib/train.py b/rllib/train.py index 691ed8fd2..b19268dfd 100755 --- a/rllib/train.py +++ b/rllib/train.py @@ -180,22 +180,23 @@ def run(args, parser): parser.error("the following arguments are required: --run") if not exp.get("env") and not exp.get("config", {}).get("env"): parser.error("the following arguments are required: --env") - if args.eager: - exp["config"]["framework"] = "tfe" - elif args.torch: + + if args.torch: exp["config"]["framework"] = "torch" - else: - exp["config"]["framework"] = "tf" + elif args.eager: + exp["config"]["framework"] = "tfe" + + if args.trace: + if exp["config"]["framework"] not in ["tf2", "tfe"]: + raise ValueError("Must enable --eager to enable tracing.") + exp["config"]["eager_tracing"] = True + if args.v: exp["config"]["log_level"] = "INFO" verbose = 2 if args.vv: exp["config"]["log_level"] = "DEBUG" verbose = 3 - if args.trace: - if exp["config"]["framework"] != "tfe": - raise ValueError("Must enable --eager to enable tracing.") - exp["config"]["eager_tracing"] = True if args.ray_num_nodes: cluster = Cluster()