[BUG] Error: Sizes of tensors must match except in dimension 1. Expected size 64 but got size 512 for tensor number 2 in the list.

Describe the bug
When running an inference request,

To Reproduce
Steps to reproduce the behavior:

1. Mask and image size of 512
2. What packages are required and their versions

python = ">=3.10,<3.12"
click = ">=7,<8"
supabase = "^2.3.4"
distro-info = "1.0"
scikit-build = "^0.16.7"
pillow = "^10.2.0"
ninja = "^1.11.1"
diffusers = "^0.26.2"
# optimization of stable-diffusion?
triton = "^2.2.0"
transformers = "^4.37.2"
pynsq = "^0.9.1"
compel = "^2.0.2"
scipy = "^1.12.0"
torch = "^2.2.0"
xformers = "^0.0.24"
accelerate = "^0.26.1"
deepspeed = { git = "https://github.com/microsoft/DeepSpeed.git", branch = "lekurile/fix_sd_ci" }

4. How to run the script

    with torch.inference_mode():
        pipe = StableDiffusionInpaintPipeline.from_pretrained(
            "/home/mikob/tools/stable-diffusion-v1-5-inpainting",
            # "/app/stable-diffusion-inpainting",
            # 'runwayml/stable-diffusion-inpainting'
            revision="fp16",
            torch_dtype=torch.float16,
            safety_checker=None,
        )

        pipe = pipe.to("cuda")
        pipe.safety_checker = dummy_safety_checker
        compel = Compel(tokenizer=pipe.tokenizer, text_encoder=pipe.text_encoder)

        # memory savings
        # pipe.enable_sequential_cpu_offload()
        # memory savings
        # pipe.enable_attention_slicing(1)
        # doesn't seem to have an effect
        # torch.backends.cuda.matmul.allow_tf32 = True
        # pipe.enable_xformers_memory_efficient_attention()

        pipe = deepspeed.init_inference(
            model=getattr(pipe,"model", pipe),
            mp_size=torch.cuda.device_count(),        # Number of GPUs
            dtype=torch.half, # dtype of the weights (fp16)
            replace_with_kernel_inject=True, # replace the model with the kernel injector
            # enable_cuda_graph=True, # enable the CUDA graph
        )

5. ...

Expected behavior
My inferences run fine without deepspeed, expect them to run the same way with it.

Screenshots
If applicable, add screenshots to help explain your problem.

System info (please complete the following information):

• OS: Linux Mint 21.2
• GPU count and types 3070ti laptop 8gb

Docker context
Are you using a specific docker image that you can share?

Additional context
Add any other context about the problem here.