From 369a3417c44c3fb0e1d6b30146d7be3f8978e1d8 Mon Sep 17 00:00:00 2001
From: Sven Mika <sven@anyscale.io>
Date: Fri, 27 Mar 2020 06:07:10 +0100
Subject: [PATCH] [RLlib] Add tf-graph by default when doing
 `Policy.export_model()`. (#7759)

* Rollback.

* WIP.

* WIP.

* Fix.

* LINT.
---
 rllib/policy/tf_policy.py    | 4 +++-
 rllib/policy/torch_policy.py | 1 -
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/rllib/policy/tf_policy.py b/rllib/policy/tf_policy.py
index 779fa303b..d6310ae88 100644
--- a/rllib/policy/tf_policy.py
+++ b/rllib/policy/tf_policy.py
@@ -347,7 +347,9 @@ class TFPolicy(Policy):
             signature_def_map = self._build_signature_def()
             builder.add_meta_graph_and_variables(
                 self._sess, [tf.saved_model.tag_constants.SERVING],
-                signature_def_map=signature_def_map)
+                signature_def_map=signature_def_map,
+                saver=tf.summary.FileWriter(export_dir).add_graph(
+                    graph=self._sess.graph))
             builder.save()
 
     @override(Policy)
diff --git a/rllib/policy/torch_policy.py b/rllib/policy/torch_policy.py
index eed61f42f..4e51f70e9 100644
--- a/rllib/policy/torch_policy.py
+++ b/rllib/policy/torch_policy.py
@@ -233,7 +233,6 @@ class TorchPolicy(Policy):
 
     def extra_grad_info(self, train_batch):
         """Return dict of extra grad info."""
-
         return {}
 
     def optimizer(self):