get_layer_output.py

def get_layer_output(model, layer, x):
    layer_output = None
    def layer_output_hook(m, i, o):
        layer_output = o.clone()
    hook = layer.register_forward_hook(layer_output_hook)
    _ = model(x)  # call forward hook
    hook.remove()
    return layer_output