TCAV/ActivationMaximization gets wrong size

I am trying to calculate the Testing for Concept Activation Vector for a certain layer and for multiple concepts. To do this, I am using the code from: iMIMIC-RCVs/rcv_utils.py at master · maragraziani/iMIMIC-RCVs · GitHub

def compute_tcav(model, layer_idx, filter_indices, seed_input,
                       wrt_tensor=None, backprop_modifier=None, grad_modifier='absolute'):
    if backprop_modifier is not None:
        modifier_fn = get(backprop_modifier)
        model = modifier_fn(model)

    # `ActivationMaximization` loss reduces as outputs get large, hence negative gradients indicate the direction
    # for increasing activations. Multiply with -1 so that positive gradients indicate increase instead.
    losses = [
        (ActivationMaximization(model.layers[layer_idx], filter_indices), -1)
    ]
    
    opt = Optimizer(input_tensor, losses, wrt_tensor=wrt_tensor, norm_grads=False)
    grads = opt.minimize(seed_input=seed_input, max_iter=1, grad_modifier=grad_modifier, verbose=False)[1]

    channel_idx = 1 if K.image_data_format() == 'channels_first' else -1
    #grads = np.max(grads, axis=channel_idx)
    return utils.normalize(grads)[0]

When the code above is used, it is used like:

nnn=rcv_utils.compute_tcav(bc_model,-1,0, np.expand_dims(test_inputs[p], axis=0), wrt_tensor=bc_model.get_layer('res4a').output)

As all my models are in Pytorch, I “translated” the code above from Pytorch to Keras and I now have the following:


from torch.autograd import Variable
def compute_tcav(model, layer_name, main_dir, device='cuda'):
    d
ataset = OLYMPOSDataset(main_dir)
    loader = torch.utils.data.DataLoader(dataset, batch_size=1)
    optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
    
    # Function to retrieve activation
    activation = {}
    tcav = {}
    def hook_fn(m, i, o):
        activation[m] = o#.detach() 
    
    
    for name, layer in model.named_modules(): 
        if name.startswith(layer_name):
            print("Hook made for...", name)
            layer.register_forward_hook(hook_fn) 
            L = layer

    model.eval()
    
    tcav = {}
    for ind, (img, label) in enumerate(loader):
        img.requires_grad=True
        img = img.to(device, dtype=torch.float)
        
        output = model(img)
        
        layer_activation =  activation[L]
        
        loss = -torch.mean(layer_activation)
        loss.backward(retain_graph=True)
        grad = tuple_of_tensors_to_tensor(torch.autograd.grad(loss, img))
        grads = normalize(grad.cpu().numpy().ravel())
        tcav[ind] = grads
    
    return tcav

However, I later need to calculate the dot product between the TCAV and the regression concept vector (RCV, not show in this question). The size of RCV is the same size as the activation array of a specific layer.

My question is, what needs to be changed in my code in order to get an output of compute_tcav() that is the same size as the activation array in question? Then I would be able to calculate the dot product.