[BUG] Initialization change for the dense layer in the TF implem of ResNet18 is not applied correctly

[zaccharieramzi]

Indeed, because I am changing the model itself, iterating through the layers doesn't give me the last dense layer.

The change is therefore never applied when I remove the initial downsampling.

When I do not, there is the following bug:

Traceback (most recent call last):
  File "/gpfsdswork/projects/rech/xpa/uap69lx/benchopt/benchopt/utils/pdb_helpers.py", line 28, in exception_handler
    yield ctx
  File "/gpfsdswork/projects/rech/xpa/uap69lx/benchopt/benchopt/runner.py", line 103, in run_one_to_cvg
    solver.run(run_once_cb)
  File "/gpfsdswork/projects/rech/xpa/uap69lx/benchmark_resnet_classif/utils/tf_solver.py", line 167, in run
    self.model = self.model_init_fn()
  File "/gpfsdswork/projects/rech/xpa/uap69lx/benchmark_resnet_classif/objective.py", line 120, in _model_init_fn
    model = model_klass(
  File "/gpfsdswork/projects/rech/xpa/uap69lx/benchmark_resnet_classif/utils/tf_resnets.py", line 331, in ResNet18
    change_dense_init(model)
  File "/gpfsdswork/projects/rech/xpa/uap69lx/benchmark_resnet_classif/utils/tf_resnets.py", line 264, in change_dense_init
    layer.build(layer.input_spec.shape)
  File "/gpfswork/rech/xpa/uap69lx/.local_torch/lib/python3.9/site-packages/keras/layers/core/dense.py", line 141, in build
    raise ValueError('The last dimension of the inputs to a Dense layer '
ValueError: The last dimension of the inputs to a Dense layer should be defined. Found None. Full input shape received: <unknown>

This is because the input specifications are not defined correctly despite the layers being built, I dk why.
I could simply base this shape on the current kernel shape.