Merge pull request #39717 from SwapnilDreams100:master
PiperOrigin-RevId: 312847670 Change-Id: I7f2d4a5cf63b55ce6d013ceb49a30ab342efa7bf
This commit is contained in:
commit
19f4ac7694
|
@ -897,7 +897,7 @@ class TPUExtended(distribute_lib.StrategyExtendedV1):
|
||||||
if tensor_util.is_tensor(input_tensor):
|
if tensor_util.is_tensor(input_tensor):
|
||||||
rank = input_tensor.get_shape().rank
|
rank = input_tensor.get_shape().rank
|
||||||
else:
|
else:
|
||||||
rank = np.rank(input_tensor)
|
rank = np.ndim(input_tensor)
|
||||||
maximum_shape = tensor_shape.TensorShape([None] * rank)
|
maximum_shape = tensor_shape.TensorShape([None] * rank)
|
||||||
maximum_shapes.append(maximum_shape)
|
maximum_shapes.append(maximum_shape)
|
||||||
maximum_shapes = nest.pack_sequence_as(replicate_inputs[0],
|
maximum_shapes = nest.pack_sequence_as(replicate_inputs[0],
|
||||||
|
|
Loading…
Reference in New Issue