Mismatch error when I try to inference with small models

Hi, when I run small-small models to inpainting with default images provided I encounter this mismatch error:

Traceback (most recent call last):
  File "/home/marco/StableCascade/cn.py", line 104, in <module>
    for (sampled_c, _, _) in tqdm(sampling_c, total=extras.sampling_configs['timesteps']):
  File "/home/marco/miniconda3/envs/casc/lib/python3.9/site-packages/tqdm/std.py", line 1181, in __iter__
    for obj in iterable:
  File "/home/marco/StableCascade/gdf/__init__.py", line 71, in sample
    pred, pred_unconditional = model(torch.cat([x, x], dim=0), noise_cond.repeat(2), **model_inputs).chunk(2)
  File "/home/marco/miniconda3/envs/casc/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1518, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/marco/miniconda3/envs/casc/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1527, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/marco/StableCascade/modules/stage_c.py", line 246, in forward
    level_outputs = self._down_encode(x, r_embed, clip, cnet)
  File "/home/marco/StableCascade/modules/stage_c.py", line 182, in _down_encode
    x = x + nn.functional.interpolate(next_cnet, size=x.shape[-2:], mode='bilinear',
RuntimeError: The size of tensor a (1536) must match the size of tensor b (2048) at non-singleton dimension 1

I also printed the two tensors involved in this error (stage_c.py):

torch.Size([8, 1536, 24, 24])
torch.Size([8, 2048, 24, 24])

Anyone with the same problem? Maybe it's my fault but I don't know how to solve it. Thanks

Stability-AI / StableCascade

Mismatch error when I try to inference with small models #51