Convert Flatten layer from PyTorch to Tensorflow - Equivalent for start_dim and end_dim

I don't think there is an identical implementation in tf. However, you can always use tf.reshape and add the shape yourself with a simple function which takes as arguments input, start_dim and end_dim and outputs the corresponding output shape that torch.flatten would give you.