Changes x.numpy().squeeze(0) to tf.squeeze(x, 0).numpy() as per the suggestion in
@@ -173,7 +173,7 @@ def to_numpy(nest: types.NestedTensor) -> types.NestedArray:
def to_numpy_squeeze(nest: types.NestedTensor, axis=0) -> types.NestedArray:
"""Converts a nest of Tensors to a nest of numpy arrays and squeeze axis."""
- return tree.map_structure(lambda x: x.numpy().squeeze(axis=axis), nest)
+ return tree.map_structure(lambda x: tf.squeeze(x, axis=axis).numpy(), nest)
def zeros_like(nest: types.Nest) -> types.NestedTensor:
