import math


def repeat_channels(conv_weights, new_channels, old_channels=3):
    """Repeat channels to match new number of input channels
    Args:
        conv_weights (torch.Tensor): shape [*, old_channels, *, *]
        new_channels (int): desired number of channels
        old_channels (int): original number of channels
    """
    rep_times = math.ceil(new_channels / old_channels)
    new_weights = conv_weights.repeat(1, rep_times, 1, 1)[:, :new_channels, :, :]
    new_weights *= old_channels / new_channels  # to keep the same output amplitude
    return new_weights