[BUG] Bloom inference error with dtype=int8

Describe the bug
When inference bloom model with replace_with_kernel_inject = True , and dtype = torch.int8

For the reason that this model is trained by torch , I load the weight with torch.load , and then use weights loaded model to init engine ( is this right ? I tried to pass checkpoint in init_inference() , but it failed )

ckpt = torch.load(self.opt.model_file, map_location='cpu') self.model.load_state_dict(ckpt['model'])

inference init :
engine = deepspeed.init_inference(model.model, mp_size = 1, dtype = torch.int8, replace_with_kernel_inject = True)

inference error :
File "/usr/local/lib/python3.7/dist-packages/deepspeed/ops/transformer/inference/ds_attention.py", line 202, in compute_attention mixed_x_layer = mixed_x_layer.view(*new_tensor_shape) RuntimeError: shape '[9, 22, 32, 240]' is invalid for input of size 506880

but with dtype = torch.half , inference success.

ds_report output

Screenshots

System info (please complete the following information):
OS: Debian GNU/Linux 10
GPU: NVIDIA A10 * 1
python : Python 3.7.3

Additional context
Question : how to load weights in init_inference() with weights generated by torch.save() ?