def merge_first_two_dims(tensor):
shape = tensor.shape<EOL>batch_size = shape[<NUM_LIT:0>] * shape[<NUM_LIT:1>]<EOL>new_shape = tuple([batch_size] + list(shape[<NUM_LIT:2>:]))<EOL>return tensor.view(new_shape)<EOL>
Reshape tensor to merge first two dimensions