def METHOD_NAME(op: AbstractPrimitive) -> AbstractPrimitive:
    """ Makes batchnorms in the op affine, if they exist """
    init_params = op.init_params
    has_batchnorm = False
    for module in op.modules():
        if isinstance(module, torch.nn.BatchNorm2d):
            has_batchnorm = True
            break
    if not has_batchnorm:
        return op
    if 'affine' in init_params:
        init_params['affine'] = True
    if 'track_running_stats' in init_params:
        init_params['track_running_stats'] = True
    new_op = type(op)(**init_params)
    return new_op