[BUG] - Channel visualization for other model

[RuoyuChen10]

Select the modules to which the bug refers:

• Attributions Methods
• Feature Visualization
• Concepts
• Metrics
• Documentation

Describe the bug
When I use another model and the last bn layer to reproduce the results of channel visualization:

obj_channels = Objective.channel(face_model, 'stack4_block3_3_bn', list(range(12)))
imgs, obj_names = optimize(obj_channels, 
                  nb_steps=1024,
                  optimizer=tf.keras.optimizers.Adam(0.05))

for i in range(len(classes)):
  plt.subplot(len(classes)//4, 4, i+1)
  plt.imshow(imgs[0][i])
  plt.title(obj_names[i])
  plt.axis('off')

it occurs bug:

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[27], line 2
      1 obj_channels = Objective.channel(model, 'stack4_block3_3_bn', list(range(12)))
----> 2 imgs, obj_names = optimize(obj_channels, 
      3                   nb_steps=1024,
      4                   optimizer=tf.keras.optimizers.Adam(0.05))
      6 for i in range(len(classes)):
      7   plt.subplot(len(classes)//4, 4, i+1)

File [~/anaconda3/envs/smdl/lib/python3.8/site-packages/xplique/features_visualizations/optim.py:112](https://vscode-remote+wsl-002bubuntu-002d22-002e04.vscode-resource.vscode-cdn.net/mnt/e/SMDL/xplique_fold/~/anaconda3/envs/smdl/lib/python3.8/site-packages/xplique/features_visualizations/optim.py:112), in optimize(objective, optimizer, nb_steps, use_fft, fft_decay, std, regularizers, image_normalizer, values_range, transformations, warmup_steps, custom_shape, save_every)
    110 images_optimized = []
    111 for step_i in range(nb_steps):
--> 112     grads = optimisation_step(model, inputs)
    113     optimizer.apply_gradients([(-grads, inputs)])
    115     last_iteration = step_i == nb_steps - 1

File [~/anaconda3/envs/smdl/lib/python3.8/site-packages/tensorflow/python/util/traceback_utils.py:153](https://vscode-remote+wsl-002bubuntu-002d22-002e04.vscode-resource.vscode-cdn.net/mnt/e/SMDL/xplique_fold/~/anaconda3/envs/smdl/lib/python3.8/site-packages/tensorflow/python/util/traceback_utils.py:153), in filter_traceback..error_handler(*args, **kwargs)
    151 except Exception as e:
    152   filtered_tb = _process_traceback_frames(e.__traceback__)
--> 153   raise e.with_traceback(filtered_tb) from None
    154 finally:
    155   del filtered_tb

File /tmp/__autograph_generated_filefwi4_a_x.py:53, in outer_factory..inner_factory..tf__step(model, inputs)
...
    File "/home/cry/anaconda3/envs/smdl/lib/python3.8/site-packages/xplique/features_visualizations/objectives.py", line 264, in optim_func  *
        return tf.reduce_mean(output * target, axis=axis_to_reduce)

    TypeError: Input 'y' of 'Mul' Op has type float32 that does not match type float16 of argument 'x'.

However, when I use it to run neuro visualization, it's ok.

I also print the model arch. :

__________________________________________________________________________________________________
 Layer (type)                Output Shape                 Param #   Connected to                  
==================================================================================================
 input_1 (InputLayer)        [(None, 112, 112, 3)]        0         []                            
                                                                                                  
 zero_padding2d (ZeroPaddin  (None, 114, 114, 3)          0         ['input_1[0][0]']             
 g2D)     


.........

 stack4_block3_2_prelu (PRe  (None, 7, 7, 512)            512       ['stack4_block3_2_bn[0][0]']  
 LU)                                                                                              
                                                                                                  
 zero_padding2d_98 (ZeroPad  (None, 9, 9, 512)            0         ['stack4_block3_2_prelu[0][0]'
 ding2D)                                                            ]                             
                                                                                                  
 stack4_block3_2_conv (Conv  (None, 7, 7, 512)            2359296   ['zero_padding2d_98[0][0]']   
 2D)                                                                                              
                                                                                                  
 stack4_block3_3_bn (BatchN  (None, 7, 7, 512)            2048      ['stack4_block3_2_conv[0][0]']
 ormalization)                                                                                    
                                                                                                  
 stack4_block3_add (Add)     (None, 7, 7, 512)            0         ['stack4_block2_add[0][0]',   
                                                                     'stack4_block3_3_bn[0][0]']  
                                                                                                  
 E_batchnorm (BatchNormaliz  (None, 7, 7, 512)            2048      ['stack4_block3_add[0][0]']   
 ation)                                                                                           
                                                                                                  
 dropout (Dropout)           (None, 7, 7, 512)            0         ['E_batchnorm[0][0]']         
                                                                                                  
 E_flatten (Flatten)         (None, 25088)                0         ['dropout[0][0]']             
                                                                                                  
 E_dense (Dense)             (None, 512)                  1284505   ['E_flatten[0][0]']           
                                                          6                                       
                                                                                                  
 pre_embedding (BatchNormal  (None, 512)                  2048      ['E_dense[0][0]']             
 ization)                                                                                         
                                                                                                  
 embedding (Activation)      (None, 512)                  0         ['pre_embedding[0][0]']       
                                                                                                  
 arcface (NormDense)         (None, 10177)                5210624   ['embedding[0][0]']           
                                                                                                  
==================================================================================================
Total params: 70431552 (268.68 MB)
Trainable params: 70362048 (268.41 MB)
Non-trainable params: 69504 (271.50 KB)
__________________________________________________________________________________________________

Expected behavior
I want to get the channel visualization results.