So this is correct:

class View(nn.Module):

    def __init__(self, shape):
        self.shape = shape

    def forward(self, input):
        '''
        TODO: the first dimension is the data batch_size
        so we need to decide how the input shape should be like
        '''
        return input.view(*self.shape)