def METHOD_NAME(self, module, input):
    if isinstance(self.weight, torch.nn.parameter.UninitializedParameter):
        x = input[0][0] if isinstance(input, tuple) else input[0]
        in_channels = x.size(-1)
        self.weight.materialize((self.K, in_channels, self.out_channels))
        size = self.weight.size(0) * self.weight.size(1)
        uniform(size, self.weight)
    module._hook.remove()
    delattr(module, '_hook')