Method register_hook

knn_utils/saveEmbedMulti.py:205–207

(self, layer, func, pre=False)

        return knn_log_probs
    
    def register_hook(self, layer, func, pre=False):
        handle = layer.register_forward_pre_hook(func) if pre else layer.register_forward_hook(func)
        self.hook_handles.append(handle)

    def break_out(self):
        for h in self.hook_handles: