How to use Tensor Parallelism for WizardCoder Inference?

I am trying to use DeepSpeed Tensor Parallelism for WizardCoder. I am making modifications in the following script. The only major change I am making is the difference in the model_name passed as WizardLM/WizardCoder-15B-V1.0. However, despite passing num_gpus as 2 (the model should easily fit into 2 GPUs, each having 48GB memory), the process hangs for about 5 mins and then exits with the below error trace.
Can someone point out what I might be doing wrong ?

Traceback (most recent call last):
  File "inference-test.py", line 126, in <module>
    outputs = pipe(inputs,
  File "/home/anmol/TieredModels/code/02_deepspeed_inference/utils.py", line 76, in __call__
    outputs = self.generate_outputs(input_list, num_tokens=num_tokens, do_sample=do_sample)
  File "/home/anmol/TieredModels/code/02_deepspeed_inference/utils.py", line 126, in generate_outputs
    outputs = self.model.generate(**input_tokens, **generate_kwargs)
  File "/home/anmol/anaconda3/envs/wizard_coder/lib/python3.8/site-packages/deepspeed/inference/engine.py", line 588, in _generate
    return self.module.generate(*inputs, **kwargs)
  File "/home/anmol/anaconda3/envs/wizard_coder/lib/python3.8/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/home/anmol/anaconda3/envs/wizard_coder/lib/python3.8/site-packages/transformers/generation/utils.py", line 1588, in generate
    return self.sample(
  File "/home/anmol/anaconda3/envs/wizard_coder/lib/python3.8/site-packages/transformers/generation/utils.py", line 2642, in sample
    outputs = self(
  File "/home/anmol/anaconda3/envs/wizard_coder/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/anmol/anaconda3/envs/wizard_coder/lib/python3.8/site-packages/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py", line 807, in forward
    transformer_outputs = self.transformer(
  File "/home/anmol/anaconda3/envs/wizard_coder/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/anmol/anaconda3/envs/wizard_coder/lib/python3.8/site-packages/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py", line 672, in forward
    outputs = block(
  File "/home/anmol/anaconda3/envs/wizard_coder/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/anmol/anaconda3/envs/wizard_coder/lib/python3.8/site-packages/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py", line 316, in forward
    attn_outputs = self.attn(
  File "/home/anmol/anaconda3/envs/wizard_coder/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/anmol/anaconda3/envs/wizard_coder/lib/python3.8/site-packages/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py", line 230, in forward
    query, key_value = self.c_attn(hidden_states).split((self.embed_dim, 2 * self.kv_dim), dim=2)
  File "/home/anmol/anaconda3/envs/wizard_coder/lib/python3.8/site-packages/torch/_tensor.py", line 803, in split
    return torch._VF.split_with_sizes(self, split_size, dim)
RuntimeError: split_with_sizes expects split_sizes to sum exactly to 3200 (input tensor's size at dimension 2), but got split_sizes=[3072, 256]
Traceback (most recent call last):
  File "inference-test.py", line 126, in <module>
    outputs = pipe(inputs,
  File "/home/anmol/TieredModels/code/02_deepspeed_inference/utils.py", line 76, in __call__
    outputs = self.generate_outputs(input_list, num_tokens=num_tokens, do_sample=do_sample)
  File "/home/anmol/TieredModels/code/02_deepspeed_inference/utils.py", line 126, in generate_outputs
    outputs = self.model.generate(**input_tokens, **generate_kwargs)
  File "/home/anmol/anaconda3/envs/wizard_coder/lib/python3.8/site-packages/deepspeed/inference/engine.py", line 588, in _generate
    return self.module.generate(*inputs, **kwargs)
  File "/home/anmol/anaconda3/envs/wizard_coder/lib/python3.8/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/home/anmol/anaconda3/envs/wizard_coder/lib/python3.8/site-packages/transformers/generation/utils.py", line 1588, in generate
    return self.sample(
  File "/home/anmol/anaconda3/envs/wizard_coder/lib/python3.8/site-packages/transformers/generation/utils.py", line 2642, in sample
    outputs = self(
  File "/home/anmol/anaconda3/envs/wizard_coder/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/anmol/anaconda3/envs/wizard_coder/lib/python3.8/site-packages/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py", line 807, in forward
    transformer_outputs = self.transformer(
  File "/home/anmol/anaconda3/envs/wizard_coder/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/anmol/anaconda3/envs/wizard_coder/lib/python3.8/site-packages/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py", line 672, in forward
    outputs = block(
  File "/home/anmol/anaconda3/envs/wizard_coder/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/anmol/anaconda3/envs/wizard_coder/lib/python3.8/site-packages/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py", line 316, in forward
    attn_outputs = self.attn(
  File "/home/anmol/anaconda3/envs/wizard_coder/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/anmol/anaconda3/envs/wizard_coder/lib/python3.8/site-packages/transformers/models/gpt_bigcode/modeling_gpt_bigcode.py", line 230, in forward
    query, key_value = self.c_attn(hidden_states).split((self.embed_dim, 2 * self.kv_dim), dim=2)
  File "/home/anmol/anaconda3/envs/wizard_coder/lib/python3.8/site-packages/torch/_tensor.py", line 803, in split
    return torch._VF.split_with_sizes(self, split_size, dim)
RuntimeError: split_with_sizes expects split_sizes to sum exactly to 3200 (input tensor's size at dimension 2), but got split_sizes=[3072, 256]

Start with inference-test.py and the referenced utils.py, then follow the traceback into the GPTBigCode attention path. Reproduce the WizardCoder run with num_gpus set to 2 and inspect the model configuration alongside the reported tensor dimensions. Done means tensor-parallel inference completes without the split_with_sizes error.