def privatize_model()

in modeling.py [0:0]

• 16 lines of code
• 8 McCabe index (conditional complexity)

def privatize_model(model, clip, std):
    """
    Converts a "normal" model into a model that computes private gradients.
    """
    types = tuple(getattr(nn, mod_type) for mod_type in MODULE_TYPES)
    private_types = tuple(getattr(modules, mod_type) for mod_type in MODULE_TYPES)
    for module in model.modules():
        if isinstance(module, types) and not isinstance(module, private_types):
            typename = str(type(module))
            typename = typename[typename.rfind(".") + 1:-2]
            module.__class__ = getattr(modules, typename)
            module.clip = torch.tensor(clip)
            module.std = torch.tensor(std)
        else:
            if hasattr(module, "weight") or hasattr(module, "bias"):
                raise NotImplementedError(
                    f"Privacy conversion of {type(module)} not implemented."
                )
    return model