ValueError: Trying to set a tensor of shape torch.Size([16, 1536, 1, 1]) in "weight" (which has shape torch.Size([16, 2048, 1, 1])), this look incorrect.

[loboere]

I try to load the small-small models but it generates an error

['transforms', 'clip_preprocess', 'gdf', 'sampling_configs', 'effnet_preprocess']
models/stage_c_lite_bf16.safetensors
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
[<ipython-input-4-1d475d0c6014>](https://localhost:8080/#) in <cell line: 3>()
      1 # SETUP MODELS & DATA
      2 extras = core.setup_extras_pre()
----> 3 models = core.setup_models(extras)
      4 models.generator.eval().requires_grad_(False)
      5 print("STAGE C READY")

1 frames
[/usr/local/lib/python3.10/dist-packages/accelerate/utils/modeling.py](https://localhost:8080/#) in set_module_tensor_to_device(module, tensor_name, device, value, dtype, fp16_statistics, tied_params_map)
    343     if value is not None:
    344         if old_value.shape != value.shape:
--> 345             raise ValueError(
    346                 f'Trying to set a tensor of shape {value.shape} in "{tensor_name}" (which has shape {old_value.shape}), this look incorrect.'
    347             )

ValueError: Trying to set a tensor of shape torch.Size([16, 1536, 1, 1]) in "weight" (which has shape torch.Size([16, 2048, 1, 1])), this look incorrect.

stage_c_3b.yaml

# GLOBAL STUFF
model_version: 3.6B
dtype: bfloat16

effnet_checkpoint_path: models/effnet_encoder.safetensors
previewer_checkpoint_path: models/previewer.safetensors
generator_checkpoint_path: models/stage_c_lite_bf16.safetensors

[pabloppp]

In order to use the small stage C, the field "model version" should be set to 1B in the config Yaml (and 700M for stageB)

https://github.com/Stability-AI/StableCascade/blob/209a52600f35dfe2a205daef54c0ff4068e86bc7/train/train_c.py#L152

https://github.com/Stability-AI/StableCascade/blob/209a52600f35dfe2a205daef54c0ff4068e86bc7/train/train_b.py#L167

[alistairwoodcock]

Just a note for anyone else who ran into this issue, if you're using the yaml config/inference/ files from the notebooks, you need to set model_version to 1B