def METHOD_NAME(module, input_channels_to_prune):
    """
    :param module: module
    :param input_channels_to_prune: list of input channels to be zeroed out
    :return:
    """
    for input_channel in input_channels_to_prune:
        if isinstance(module, nn.Conv2d):
            module.weight.data[:, input_channel, :, :] = 0
        elif isinstance(module, nn.Linear):
            module.weight.data[:, input_channel] = 0
        else:
            raise ValueError("Unsupported layer_type")