import torch


def get_output_size(model, shape, channels, cuda):
    """
    Get the size of the output of the last layer of the model.

    :model: the model
    :shape: shape of the input image tuple(width, height)
    :channels: amount of channels of input image (int)
    :cuda: wether or not to use GPU
    :returns: the size of the output of the last layer of the model.
    """
    bs = 1

    input = torch.rand(bs, channels, *shape)
    if cuda:
        input = input.cuda()
        model = model.cuda()
    output_feat = model.forward(input)
    n_size = output_feat.data.view(bs, -1).size(1)
    return n_size