This will break once kwargs has no key of 'partitioner'.
PiperOrigin-RevId: 310463715 Change-Id: I45fcd670172a8d96f509fbf3e4ba64556386b058
This commit is contained in:
parent
79acb0824b
commit
491d6e42ce
@ -1353,7 +1353,7 @@ def split_compile_and_replicate(computation,
|
|||||||
|
|
||||||
def custom_getter(getter, name, *args, **kwargs):
|
def custom_getter(getter, name, *args, **kwargs):
|
||||||
"""Variables on TPU have a few restrictions."""
|
"""Variables on TPU have a few restrictions."""
|
||||||
partitioner = kwargs["partitioner"]
|
partitioner = kwargs.get("partitioner", None)
|
||||||
if partitioner is not None:
|
if partitioner is not None:
|
||||||
kwargs["partitioner"] = None
|
kwargs["partitioner"] = None
|
||||||
logging.warning(
|
logging.warning(
|
||||||
|
Loading…
x
Reference in New Issue
Block a user