INNER CODE UNIT · Python

hook

google/fully-homomorphic-encryption · demos/cc_fraud/cleartext/estimate_ranges.py:98

    def hook(module, input_tensor, output):
      activation_outputs[name] = output.detach()

    # pylint: enable=unused-argument

    return hook

  hooks = []
  linear_count = 0
  for name, module in model.named_modules():
    if isinstance(module, nn.Linear) and module is not model.net[-1]:
      hook_name = f"Linear_{linear_count}"
      hooks.append(module.register_forward_hook(make_hook(hook_name)))
      linear_count += 1

  print(f"Registered {len(hooks)} hooks on Linear layers.")

  # Run evaluation

View source record →