This will break once kwargs has no key of 'partitioner'.

PiperOrigin-RevId: 310463715 Change-Id: I45fcd670172a8d96f509fbf3e4ba64556386b058
2020-05-07 16:56:39 -07:00 · 2020-05-07 16:56:39 -07:00 · 491d6e42ce
commit 491d6e42ce
parent 79acb0824b
1 changed files with 1 additions and 1 deletions
--- a/tensorflow/python/tpu/tpu.py
+++ b/tensorflow/python/tpu/tpu.py
@ -1353,7 +1353,7 @@ def split_compile_and_replicate(computation,

      def custom_getter(getter, name, *args, **kwargs):
        """Variables on TPU have a few restrictions."""
-        partitioner = kwargs["partitioner"]
+        partitioner = kwargs.get("partitioner", None)
        if partitioner is not None:
          kwargs["partitioner"] = None
          logging.warning(