I’m implementing a custom layer that needs multiple reshaping of two input tensors, THE BATCH IS  OF THE SHAPE R^(N, a,b,c), where N is the number of tensors in the batch. Actually I need to reshape each 3D tensor with the size of R^(a,b,c) to be a matrix of the size (a*b,c) and then use torch.mm to multiply them. But I keep running into an error for the inconsistency of size (because torch.mm requires the inputs to be of 2D shape rather than 3D (in case of N>1) while custom layer reads the input as a four dimensional tensor (N,a,b,c). Any help will be highly appreciated.