Update tpu_strategy.py

rank[DEP] -> ndim
This commit is contained in:
Swapnil Parekh 2020-05-20 19:13:15 +05:30 committed by GitHub
parent cdd4e5e918
commit 18a7a327f3
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -897,7 +897,7 @@ class TPUExtended(distribute_lib.StrategyExtendedV1):
if tensor_util.is_tensor(input_tensor):
rank = input_tensor.get_shape().rank
else:
rank = np.rank(input_tensor)
rank = np.ndim(input_tensor)
maximum_shape = tensor_shape.TensorShape([None] * rank)
maximum_shapes.append(maximum_shape)
maximum_shapes = nest.pack_sequence_as(replicate_inputs[0],